"""How exploitable is a frozen policy?

Beating an opponent on average and being hard to beat are different properties.
A policy can win the head-to-head and still have a habit a dedicated opponent can
farm -- rock-paper-scissors throwing rock 60% of the time beats a lot of people
and loses every time to anyone who noticed.

So: freeze the policy, train a fresh one from scratch whose only job is to beat
*that* policy, and see how far it gets. Near 0.5 means there was nothing to find.
Well above it means there was.

PPO self-play carries no Nash guarantee, so this is not a formality. The number it
produces is a lower bound on exploitability -- a stronger exploiter might find
more -- so it is only meaningful compared against the same exploiter budget spent
on another policy.
"""

from __future__ import annotations

import json
import time
from collections.abc import Callable
from pathlib import Path

import jax
import jax.numpy as jnp

from lost_cities_jax.match import MatchState, match_legal_action_mask, match_reset, match_step
from lost_cities_jax.match_ppo import (
    FULL,
    N_SEATS,
    Ablation,
    MatchActorCritic,
    MatchTransition,
    _seat_views,
    create_match_train_state,
    match_lead,
    match_ppo_update,
)
from lost_cities_jax.obs import observation
from lost_cities_jax.ppo import (
    ActorCritic,
    JaxPPOConfig,
    TrainState,
    _append_jsonl,
    _broadcast_done,
    action_log_prob,
    categorical_entropy,
    compute_gae,
    mask_logits,
    masked_mean,
)
from lost_cities_jax.types import PLAY

FrozenPolicy = Callable[[MatchState, jax.Array], jax.Array]
"""(env batch, seat batch) -> greedy actions. Whatever architecture, hidden here."""


def match_frozen_policy(cfg: JaxPPOConfig, params, ablation: Ablation = FULL) -> FrozenPolicy:
    model = MatchActorCritic(cfg.network.hidden_size, cfg.network.num_layers)
    seats = jnp.arange(N_SEATS, dtype=jnp.int32)

    def policy(env: MatchState, seat: jax.Array) -> jax.Array:
        obs, critic = _seat_views(env, seats, ablation)
        idx = seat[None, :, None]
        logits, _ = model.apply(
            params,
            jnp.take_along_axis(obs, idx, axis=0)[0],
            jnp.take_along_axis(critic, idx, axis=0)[0],
        )
        mask = jax.vmap(match_legal_action_mask)(env)
        return jnp.argmax(mask_logits(logits, mask), axis=-1).astype(jnp.int32)

    return policy


def single_round_frozen_policy(cfg: JaxPPOConfig, params) -> FrozenPolicy:
    """The old stack: it sees one round at a time and has never heard of carry."""
    model = ActorCritic(cfg.network.hidden_size, cfg.network.num_layers)

    def policy(env: MatchState, seat: jax.Array) -> jax.Array:
        logits, _ = model.apply(params, jax.vmap(observation)(env.round, seat))
        mask = jax.vmap(match_legal_action_mask)(env)
        return jnp.argmax(mask_logits(logits, mask), axis=-1).astype(jnp.int32)

    return policy


def _reset_done(env: MatchState, seat: jax.Array, rng: jax.Array, batch: int):
    env_key, seat_key = jax.random.split(rng)
    fresh = jax.vmap(match_reset)(jax.random.split(env_key, batch))
    next_env = jax.tree_util.tree_map(
        lambda old, new: jnp.where(_broadcast_done(env.done, old), new, old), env, fresh
    )
    # Redraw the exploiter's seat too, so it learns to beat the frozen policy from
    # either side rather than only the one it started in.
    fresh_seat = jax.random.bernoulli(seat_key, 0.5, shape=(batch,)).astype(jnp.int32)
    return next_env, jnp.where(env.done, fresh_seat, seat)


def make_exploiter_rollout_fn(cfg: JaxPPOConfig, frozen: FrozenPolicy):
    batch = cfg.ppo.batch_games
    seats = jnp.arange(N_SEATS, dtype=jnp.int32)
    ablation = FULL

    @jax.jit
    def rollout_fn(state: TrainState, env: MatchState, seat: jax.Array, rng: jax.Array):
        def body(carry, _):
            env, seat, key = carry
            key, act_key, reset_key = jax.random.split(key, 3)

            obs, critic_obs = _seat_views(env, seats, ablation)
            idx = seat[None, :, None]
            my_obs = jnp.take_along_axis(obs, idx, axis=0)[0]
            my_critic = jnp.take_along_axis(critic_obs, idx, axis=0)[0]
            legal = jax.vmap(match_legal_action_mask)(env)

            logits, value = state.apply_fn(state.params, my_obs, my_critic)
            masked = mask_logits(logits, legal)
            mine = jax.random.categorical(act_key, masked, axis=-1).astype(jnp.int32)
            theirs = jax.lax.stop_gradient(frozen(env, 1 - seat))

            to_move = env.round.to_move.astype(jnp.int32)
            active = ~env.done
            my_turn = to_move == seat
            actions = jnp.where(my_turn, mine, theirs)

            before = jax.vmap(match_lead)(env)
            next_env, _, _ = jax.vmap(match_step, in_axes=(0, 0))(env, actions)
            after = jax.vmap(match_lead)(next_env)

            # match_lead is seat 0's view; flip it when the exploiter sits in 1.
            sign = jnp.where(seat == 0, 1.0, -1.0)
            reward = jnp.where(active, sign * (after - before) / cfg.reward.terminal_scale, 0.0)

            actor_mask = active & my_turn
            transition = MatchTransition(
                obs=my_obs,
                critic_obs=my_critic,
                legal_mask=legal,
                action=actions,
                log_prob=action_log_prob(masked, mine),
                value=value,
                reward=reward,
                done=next_env.done,
                active=active,
                actor_mask=actor_mask,
                entropy=categorical_entropy(masked),
                play_action=((actions % 12) // 6 == PLAY) & actor_mask,
            )

            done = next_env.done
            final_lead = sign * jax.vmap(match_lead)(next_env)
            won = (done & active & (final_lead > 0)).astype(jnp.float32)
            finished = (done & active).astype(jnp.float32)

            next_env, seat = _reset_done(next_env, seat, reset_key, batch)
            return (next_env, seat, key), (transition, won, finished)

        (next_env, seat, rng), (transitions, won, finished) = jax.lax.scan(
            body, (env, seat, rng), xs=None, length=cfg.ppo.rollout_steps
        )

        obs, critic_obs = _seat_views(next_env, seats, ablation)
        idx = seat[None, :, None]
        _, last_value = state.apply_fn(
            state.params,
            jnp.take_along_axis(obs, idx, axis=0)[0],
            jnp.take_along_axis(critic_obs, idx, axis=0)[0],
        )

        played = jnp.sum(finished)
        metrics = {
            "exploiter_win_rate": jnp.sum(won) / jnp.maximum(played, 1.0),
            "matches_completed": played,
            "learner_actions": jnp.sum(transitions.actor_mask),
            "entropy_mean": masked_mean(
                transitions.entropy, transitions.actor_mask.astype(jnp.float32)
            ),
        }
        return next_env, seat, transitions, last_value, metrics

    return rollout_fn


def train_exploiter(cfg: JaxPPOConfig, frozen: FrozenPolicy, run_dir: Path) -> TrainState:
    """Train a fresh policy whose only goal is to beat ``frozen``."""

    rng = jax.random.PRNGKey(cfg.run.seed)
    rng, init_key, reset_key, seat_key = jax.random.split(rng, 4)
    state = create_match_train_state(cfg, init_key, Ablation())
    env = jax.jit(jax.vmap(match_reset))(jax.random.split(reset_key, cfg.ppo.batch_games))
    seat = jax.random.bernoulli(seat_key, 0.5, shape=(cfg.ppo.batch_games,)).astype(jnp.int32)

    rollout_fn = make_exploiter_rollout_fn(cfg, frozen)

    @jax.jit
    def iteration(state, env, seat, rng):
        rng, roll_key, update_key = jax.random.split(rng, 3)
        env, seat, transitions, last_value, metrics = rollout_fn(state, env, seat, roll_key)
        advantages, returns = compute_gae(
            transitions.reward,
            transitions.value,
            transitions.done,
            cfg.ppo.gamma,
            cfg.ppo.gae_lambda,
            last_value,
        )
        state, update_metrics = match_ppo_update(
            state, transitions, advantages, returns, update_key, cfg
        )
        return state, env, seat, rng, {**metrics, **update_metrics}

    run_dir.mkdir(parents=True, exist_ok=True)
    metrics_path = run_dir / "metrics.jsonl"
    start = time.perf_counter()
    for update in range(cfg.run.total_updates):
        state, env, seat, rng, metrics = iteration(state, env, seat, rng)
        jax.tree_util.tree_leaves(metrics)[0].block_until_ready()
        row = {k: float(v) for k, v in metrics.items()}
        row["update"] = update
        row["elapsed_seconds"] = time.perf_counter() - start
        _append_jsonl(metrics_path, row)
        if update % cfg.run.log_every == 0:
            print(
                json.dumps(
                    {
                        "update": update,
                        "exploiter_win_rate": row["exploiter_win_rate"],
                        "matches": row["matches_completed"],
                    },
                    sort_keys=True,
                ),
                flush=True,
            )
    return state


__all__ = [
    "make_exploiter_rollout_fn",
    "match_frozen_policy",
    "single_round_frozen_policy",
    "train_exploiter",
]
