#!/usr/bin/env python3
"""What is each piece of the match stack worth?

Train one variant per switch at identical compute, then play each against the full
stack in duplicate matches -- same three deals, same coins, both seats -- so deal
luck cancels and only the policy difference is left.

The variants do not share an observation shape (dropping the privileged critic or
the match features changes the input dims), so this carries its own head-to-head
rather than reusing match_eval's, which assumes one architecture.
"""

from __future__ import annotations

import json
from pathlib import Path

import jax
import jax.numpy as jnp
import numpy as np

from lost_cities_jax.match import (
    match_legal_action_mask,
    match_reset_from,
    match_score,
    match_step,
)
from lost_cities_jax.match_eval import MATCH_SCAN_STEPS, _wilson, match_bank
from lost_cities_jax.match_ppo import (
    Ablation,
    MatchActorCritic,
    _seat_views,
    create_match_train_state,
    match_train,
)
from lost_cities_jax.ppo import load_config, mask_logits, restore_checkpoint

VARIANTS = [
    Ablation(privileged_critic=False),
    Ablation(both_seats=False),
    Ablation(match_obs=False),
]
MATCHES = 4096
FULL_CHECKPOINT = Path("runs/jax-ppo-match/2026-07-15_030152_match-linear/latest")


def _load(cfg, checkpoint: Path, ablation: Ablation):
    state = create_match_train_state(cfg, jax.random.PRNGKey(0), ablation)
    return restore_checkpoint(checkpoint, state).params


def _existing_run(label: str) -> Path | None:
    runs = sorted(Path("runs/jax-ppo-match").glob(f"*ablate-{label}"))
    for run in reversed(runs):
        if (run / "latest").exists():
            return run
    return None


def duel(cfg, params_a, ablation_a, params_b, ablation_b, *, matches: int) -> dict:
    """A vs B over duplicate matches. Each side sees the world its own way."""

    decks, coins = match_bank(20260719, matches)
    model = MatchActorCritic(cfg.network.hidden_size, cfg.network.num_layers)
    seats = jnp.arange(2, dtype=jnp.int32)

    @jax.jit
    def run(env, a_seat):
        def body(carry, _):
            env, _unused = carry
            to_move = env.round.to_move.astype(jnp.int32)
            mask = jax.vmap(match_legal_action_mask)(env)

            def act(params, ablation):
                obs, critic = _seat_views(env, seats, ablation)
                idx = to_move[None, :, None]
                seat_obs = jnp.take_along_axis(obs, idx, axis=0)[0]
                seat_critic = jnp.take_along_axis(critic, idx, axis=0)[0]
                logits, _ = model.apply(params, seat_obs, seat_critic)
                return jnp.argmax(mask_logits(logits, mask), axis=-1).astype(jnp.int32)

            action = jnp.where(
                to_move == a_seat, act(params_a, ablation_a), act(params_b, ablation_b)
            )
            env, _, _ = jax.vmap(match_step, in_axes=(0, 0))(env, action)
            return (env, _unused), 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 {
        "matches": games,
        "a_win_rate": wins / games,
        "wilson_low": low,
        "wilson_high": high,
        "a_mean_lead": float(lead.mean()),
    }


def main() -> None:
    cfg = load_config("configs/jax_ppo/match-selfplay.yaml")
    full = _load(cfg, FULL_CHECKPOINT, Ablation())

    rows = []
    for ablation in VARIANTS:
        label = ablation.label()
        run_dir = _existing_run(label)
        if run_dir is None:
            print(f"\n===== training {label} =====", flush=True)
            cfg = load_config("configs/jax_ppo/match-selfplay.yaml")
            cfg.run.experiment_name = f"ablate-{label}"
            run_dir = match_train(cfg, ablation=ablation)
        else:
            print(f"\n===== reusing {run_dir} =====", flush=True)

        cfg = load_config("configs/jax_ppo/match-selfplay.yaml")
        params = _load(cfg, run_dir / "latest", ablation)
        # The ablated net plays seat A, the full stack answers. Below 0.5 means
        # the piece we removed was carrying weight.
        result = duel(cfg, params, ablation, full, Ablation(), matches=MATCHES)
        rows.append({"ablation": label, **result})
        print(f"  {label}: win rate vs full = {result['a_win_rate']:.4f}", flush=True)

    print("\n\n==================== ablation ====================")
    print(f"{'removed':<26}{'win rate vs full':>18}{'95% CI':>22}{'mean lead':>12}")
    print("-" * 78)
    for row in rows:
        ci = f"[{row['wilson_low']:.3f}, {row['wilson_high']:.3f}]"
        print(
            f"{row['ablation']:<26}{row['a_win_rate']:>18.4f}{ci:>22}{row['a_mean_lead']:>+12.1f}"
        )
    print("\nbelow 0.5 = the removed piece was carrying weight")
    Path("runs/jax-ppo-match/ablation.json").write_text(json.dumps(rows, indent=2, sort_keys=True))


if __name__ == "__main__":
    main()