#!/usr/bin/env python3
"""Freeze each policy, train an exploiter against it on the same budget, compare.

Winning the head-to-head says a policy is strong on average. It does not say the
policy is hard to beat. This does: whatever the exploiter reaches is a habit the
frozen policy could not defend.
"""

from __future__ import annotations

import argparse
import json
from pathlib import Path

import jax
import numpy as np

from lost_cities_jax.exploit import (
    match_frozen_policy,
    single_round_frozen_policy,
    train_exploiter,
)
from lost_cities_jax.match import match_reset_from, match_score
from lost_cities_jax.match_eval import MATCH_SCAN_STEPS, _wilson, match_bank
from lost_cities_jax.match_ppo import Ablation, create_match_train_state
from lost_cities_jax.ppo import create_train_state, load_config, restore_checkpoint

OURS = Path("runs/jax-ppo-match/2026-07-15_031529_match-scaled/latest")
LEAGUE = Path(
    "/mnt/2tbhdd/coolrl-lost-cities-artifacts/league/2026-07-05_052325_jax-ppo-league-v1/latest"
)
MATCHES = 4096


def _final_score(cfg, exploiter_params, frozen, matches: int) -> dict:
    """Play the trained exploiter against the frozen policy, both seats."""
    import jax.numpy as jnp

    from lost_cities_jax.exploit import match_frozen_policy as _mk
    from lost_cities_jax.match import match_step

    attacker = _mk(cfg, exploiter_params, Ablation())
    decks, coins = match_bank(20260722, matches)

    @jax.jit
    def run(env, a_seat):
        def body(carry, _):
            env, _u = carry
            action = jnp.where(
                env.round.to_move.astype(jnp.int32) == a_seat,
                attacker(env, a_seat),
                frozen(env, 1 - a_seat),
            )
            env, _, _ = jax.vmap(match_step, in_axes=(0, 0))(env, action)
            return (env, _u), None

        (env, _), _ = jax.lax.scan(body, (env, jnp.int32(0)), xs=None, length=MATCH_SCAN_STEPS)
        return env

    leads = []
    for seat in (0, 1):
        env = jax.vmap(match_reset_from)(decks, coins)
        final = run(env, jnp.full((matches,), seat, dtype=jnp.int32))
        totals = np.asarray(jax.vmap(match_score)(final))
        leads.append(totals[:, seat] - totals[:, 1 - seat])

    lead = np.concatenate(leads)
    games = float(lead.size)
    wins = float((lead > 0).sum())
    low, high = _wilson(wins, games)
    return {
        "exploiter_win_rate": wins / games,
        "wilson_low": low,
        "wilson_high": high,
        "exploiter_mean_lead": float(lead.mean()),
        "matches": games,
    }


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--updates", type=int, default=250, help="exploiter training updates")
    parser.add_argument("--batch-games", type=int, default=1024)
    parser.add_argument("--tag", default="", help="suffix for the run dir, to keep runs apart")
    args = parser.parse_args()

    match_cfg = load_config("configs/jax_ppo/match-selfplay.yaml")
    old_cfg = load_config("configs/jax_ppo/balanced.yaml")

    ours = restore_checkpoint(
        OURS, create_match_train_state(match_cfg, jax.random.PRNGKey(0), Ablation())
    ).params
    league = restore_checkpoint(LEAGUE, create_train_state(old_cfg, jax.random.PRNGKey(0))).params

    targets = {
        "ours (match stack, 131M)": match_frozen_policy(match_cfg, ours, Ablation()),
        "league (web-deployed)": single_round_frozen_policy(old_cfg, league),
    }

    rows = []
    for name, frozen in targets.items():
        print(f"\n===== training an exploiter against: {name} =====", flush=True)
        cfg = load_config("configs/jax_ppo/match-selfplay.yaml")
        cfg.ppo.batch_games = args.batch_games
        cfg.run.total_updates = args.updates
        cfg.run.log_every = max(1, args.updates // 10)
        slug = name.split()[0] + args.tag
        state = train_exploiter(cfg, frozen, Path(f"runs/jax-ppo-match/exploit-{slug}"))
        result = _final_score(cfg, state.params, frozen, MATCHES)
        rows.append({"target": name, **result})
        print(f"  exploiter reached {result['exploiter_win_rate']:.4f} vs {name}", flush=True)

    print("\n\n============ exploitability (same exploiter budget) ============")
    print(f"{'frozen policy':<28}{'exploiter win rate':>20}{'95% CI':>22}")
    print("-" * 72)
    for row in rows:
        ci = f"[{row['wilson_low']:.3f}, {row['wilson_high']:.3f}]"
        print(f"{row['target']:<28}{row['exploiter_win_rate']:>20.4f}{ci:>22}")
    print("\nhigher = the frozen policy had more to farm. 0.5 = nothing found.")
    Path(f"runs/jax-ppo-match/exploitability{args.tag}.json").write_text(json.dumps(rows, indent=2))


if __name__ == "__main__":
    main()