"""DAgger for Snake (Ross et al., 2011).

Behaviour cloning only sees states the oracle visits; once the student makes a
mistake it lands in states it never trained on and errors compound. DAgger lets
the student play, labels *its* states with the oracle, aggregates them with the
original data and fine-tunes. Cheap (no reward engineering, the oracle is
free) and directly targets the student's own failure modes.
"""
from __future__ import annotations

import copy
import json
import shutil
from pathlib import Path

from minijev.datasets.decision_data import DecisionData
from minijev.inference.runtime import TorchDecider
from minijev.training.supervised import train

from .dataset import QUESTION, collect_with_policy


def run_dagger(cfg: dict, checkpoint: str, iterations: int = 2, samples: int = 100_000,
               run_dir: str | None = None) -> list[dict]:
    import torch

    # model/training settings come from the checkpoint (DAgger fine-tunes *that* model);
    # the dagger config only contributes its `dagger`, `env` and `eval` sections
    ck_cfg = torch.load(checkpoint, map_location="cpu", weights_only=False)["config"]
    cfg = {**copy.deepcopy(ck_cfg), **{k: cfg[k] for k in ("dagger", "env", "eval") if k in cfg}}
    device = "cuda" if torch.cuda.is_available() else "cpu"
    root = Path(run_dir or f"runs/{ck_cfg.get('name', 'snake')}-dagger")
    root.mkdir(parents=True, exist_ok=True)
    base_dir = Path(cfg["data"]["dir"])
    agg = DecisionData.load(base_dir / "train.npz")
    dcfg = cfg.get("dagger", {})
    board_sizes = cfg["data"].get("generate", {}).get("board_sizes", [cfg.get("env", {}).get("board_size", 10)])
    history = json.loads((root / "dagger.json").read_text()) if (root / "dagger.json").exists() else []
    for rec in history:  # resume: finished iterations are kept, their data is rebuilt below
        checkpoint = rec["checkpoint"]
    if history:
        agg = DecisionData.load(root / f"data-it{history[-1]['iteration']}" / "train.npz")
        print(f"[dagger] resuming after iteration {history[-1]['iteration']}")
    for it in range(len(history) + 1, iterations + 1):
        dec = TorchDecider(checkpoint, threads=4, device=device)
        # probs are aligned with ACTIONS order whatever the labels ("UP" or "move up")
        policy = lambda states: dec.probs(states, QUESTION, list(dec.task.default_choices))  # noqa: E731
        print(f"[dagger] iteration {it}: collecting {samples:,} student states ...")
        new = collect_with_policy(policy, samples, seed=1000 + it, board_sizes=board_sizes,
                                  keep_prob=dcfg.get("keep_prob", 0.2), explore_eps=dcfg.get("explore_eps", 0.02),
                                  env_cfg={k: v for k, v in cfg.get("env", {}).items() if k != "board_size"},
                                  game_id_base=10**9 * it, parallel=256)
        agg = DecisionData.concat([agg, new])
        data_dir = root / f"data-it{it}"
        data_dir.mkdir(exist_ok=True)
        agg.save(data_dir / "train.npz")
        for split in ("val", "test"):  # evaluation splits stay the original, untouched ones
            shutil.copy(base_dir / f"{split}.npz", data_dir / f"{split}.npz")
        c = copy.deepcopy(cfg)
        c["name"] = f"{cfg.get('name', 'snake')}-dagger{it}"
        c["data"] = {"dir": str(data_dir)}
        c["training"]["init_checkpoint"] = checkpoint
        c["training"]["epochs"] = dcfg.get("epochs", 1)
        c["training"]["lr"] = dcfg.get("lr", cfg["training"].get("lr", 1e-3) * 0.5)
        c["training"]["warmup_steps"] = 100
        it_dir = root / f"it{it}"
        summary = train(c, run_dir=it_dir, resume=(it_dir / "checkpoints" / "last.pt").exists())
        checkpoint = summary["final_checkpoint"]
        from .evaluate import model_policy, play_games
        games = play_games(model_policy(TorchDecider(checkpoint, threads=4, device=device)),
                           cfg.get("eval", {}).get("games", 200), cfg.get("env", {}), parallel=200)
        (root / f"it{it}" / "games.json").write_text(json.dumps(games, indent=2))
        rec = {"iteration": it, "rows": len(agg), "new_rows": len(new), "checkpoint": checkpoint,
               "test_acc": summary["test_calibrated"]["accuracy"], "test_ece": summary["test_calibrated"]["ece"],
               "avg_score": games["avg_score"], "fatal_move_rate": games["fatal_move_rate"]}
        print(f"[dagger] it {it}: avg score {games['avg_score']:.1f} | test acc {rec['test_acc']:.4f} "
              f"| fatal-move {games['fatal_move_rate']:.4%}")
        history.append(rec)
        (root / "dagger.json").write_text(json.dumps(history, indent=2))
    return history