"""Evaluate Generals.io agents by playing them against each other.

Run with `.venv/bin/python scripts/evaluate.py` or `uv run python scripts/evaluate.py`.
Override defaults with OmegaConf arguments such as `num_games=200 agent0=hunter`.

Each agent is either one of the bundled JAX agents (`random`, `expander`,
`hunter`) or a trained policy loaded from an equinox checkpoint via
`agent0=model:ppo_model.eqx`. Games are run in a single vmapped batch of
`num_games` parallel envs.
"""

from dataclasses import dataclass

import jax
import jax.numpy as jnp
import jax.random as jrandom
from generals.core import game
from generals.core.env import GeneralsEnv
from omegaconf import DictConfig, OmegaConf

from general_bots_training.opponents import make_opponent


@dataclass
class EvalConfig:
    grid_dims: tuple[int, int] = (21, 21)
    truncation: int = 1200
    num_games: int = 100
    agent0: str = "random"
    agent1: str = "expander"
    seed: int = 0


def load_config(args: list[str] | None = None) -> DictConfig:
    """Merge CLI overrides into the typed default evaluation configuration."""
    defaults = OmegaConf.structured(EvalConfig)
    return OmegaConf.merge(defaults, OmegaConf.from_cli(args))  # type: ignore


def main(config: DictConfig):
    key = jrandom.PRNGKey(config.seed)
    key, env_key = jrandom.split(key, 2)

    agent0 = make_opponent(config.agent0)
    agent1 = make_opponent(config.agent1)
    label0, label1 = config.agent0, config.agent1

    env = GeneralsEnv(grid_dims=tuple(config.grid_dims), truncation=config.truncation)
    pool, _ = env.reset(env_key)
    get_obs = game.get_full_observation if env.perfect_info else game.get_observation

    n = config.num_games
    key, init_key = jrandom.split(key)
    init_keys = jrandom.split(init_key, n)
    states = jax.vmap(env.init_state)(init_keys)
    states = states._replace(pool_idx=jnp.arange(n, dtype=states.pool_idx.dtype) % env.pool_size)

    step_env = jax.vmap(env.step, in_axes=(0, 0, None))

    def run_episode(states, key):
        """Play one full game per env; return final winner per env.

        The env auto-resets on done and overwrites `state.winner` with -1, so we
        capture the winner from `timestep.info` on the first done step per env and
        latch it for the remainder of the scan.
        """

        def body(carry, _):
            states, latched_winner, key = carry
            key, p0_key, p1_key = jrandom.split(key, 3)
            obs_p0 = jax.vmap(lambda s: get_obs(s, 0))(states)
            obs_p1 = jax.vmap(lambda s: get_obs(s, 1))(states)
            keys_p0 = jrandom.split(p0_key, n)
            keys_p1 = jrandom.split(p1_key, n)
            actions_p0 = jax.vmap(agent0.act)(obs_p0, keys_p0)
            actions_p1 = jax.vmap(agent1.act)(obs_p1, keys_p1)
            actions = jnp.stack([actions_p0, actions_p1], axis=1)
            timesteps, states = step_env(states, actions, pool)
            done = timesteps.terminated | timesteps.truncated
            # On the first done, latch the winner (truncated games stay at -1 = draw).
            new_winner = jnp.where(
                done & (latched_winner < 0), timesteps.info.winner, latched_winner
            )
            return (states, new_winner, key), done

        # Run a fixed number of steps equal to the truncation length, which
        # guarantees every env has terminated or truncated at least once.
        (states, latched_winner, key), _ = jax.lax.scan(
            body, (states, jnp.full(n, -1, dtype=jnp.int32), key), None, length=config.truncation
        )
        return latched_winner

    print("Generals.io evaluation")
    print(OmegaConf.to_yaml(config, resolve=True).rstrip())
    print(f"device: {jax.devices()[0]}")
    print(f"agent0: {label0}")
    print(f"agent1: {label1}")
    print(f"games:  {n}")
    print()
    print("warming up (jit compile)...")
    winners = run_episode(states, key)
    jax.block_until_ready(winners)
    print("warming up done\n")

    print("playing...")
    winners = run_episode(states, key)
    jax.block_until_ready(winners)

    wins0 = int(jnp.sum(winners == 0))
    wins1 = int(jnp.sum(winners == 1))
    draws = int(jnp.sum(winners < 0))

    print("\nsummary")
    print("-" * 40)
    print(f"agent0 ({label0}): {wins0:4d} wins ({wins0 / n * 100:.1f}%)")
    print(f"agent1 ({label1}): {wins1:4d} wins ({wins1 / n * 100:.1f}%)")
    print(f"draws:              {draws:4d}      ({draws / n * 100:.1f}%)")
    print("-" * 40)
    print(f"total games: {n}")


if __name__ == "__main__":
    main(load_config())