#!/usr/bin/env python3
"""Train the classic three-round agent by self-play, then measure it."""

from __future__ import annotations

import argparse
import json
from pathlib import Path

import jax

from lost_cities_jax.match_eval import carry_probe, match_evaluate
from lost_cities_jax.match_ppo import create_match_train_state, match_train
from lost_cities_jax.ppo import load_config, restore_checkpoint


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", default="configs/jax_ppo/match-selfplay.yaml")
    parser.add_argument("--eval-only", type=Path, default=None)
    parser.add_argument("--matches", type=int, default=1024)
    parser.add_argument("--set", action="append", default=[])
    args = parser.parse_args()

    cfg = load_config(args.config)
    for override in args.set:
        path, _, raw = override.partition("=")
        section, _, field = path.partition(".")
        target = getattr(cfg, section)
        current = getattr(target, field)
        value = type(current)(raw) if not isinstance(current, bool) else raw == "true"
        setattr(target, field, value)

    if args.eval_only is None:
        run_dir = match_train(cfg)
        checkpoint = run_dir / "latest"
    else:
        checkpoint = args.eval_only

    state = create_match_train_state(cfg, jax.random.PRNGKey(0))
    state = restore_checkpoint(Path(checkpoint), state)

    result = match_evaluate(cfg, state.params, matches=args.matches)
    print("\n== duplicate match eval (self-play) ==")
    print(json.dumps(result, indent=2, sort_keys=True))

    print("\n== carry probe: does round-three play react to the deficit? ==")
    rows = carry_probe(cfg, state.params, matches=args.matches // 2)
    header = f"{'carry':>7}{'win_rate':>10}{'opened':>9}{'wagers':>9}{'deck_race':>11}{'plies':>8}"
    print(header)
    for row in rows:
        print(
            f"{row['carry']:>7}{row['win_rate']:>10.3f}{row['opened_colors']:>9.2f}"
            f"{row['wagers_played']:>9.2f}{row['deck_race_rate']:>11.3f}{row['mean_plies']:>8.0f}"
        )
    Path(checkpoint).parent.joinpath("carry_probe.json").write_text(
        json.dumps({"eval": result, "probe": rows}, indent=2, sort_keys=True)
    )


if __name__ == "__main__":
    main()