"""PPO fine-tuning of a supervised MiniJEV Snake policy (Phase 2).

Rollouts: `n_envs` games stepped together, batched inference, sampled actions.
Reward comes from `env.reward` in the config (food +, death -, tiny/zero per
step). The value head is warmed up first (it was never trained by behaviour
cloning), and a KL penalty to the frozen supervised policy keeps early updates
from destroying what imitation learned. Resumable from `rl_last.pt`.
"""
from __future__ import annotations

import copy
import json
import time
from pathlib import Path

import numpy as np
import torch

from minijev.calibration.metrics import calibration_report, fit_temperature
from minijev.datasets.batching import BatchEncoder
from minijev.datasets.decision_data import DecisionData
from minijev.inference.runtime import TorchDecider
from minijev.runs import RunLogger, environment_info
from minijev.tokenizer import encode_batch
from minijev.training.ppo import gae, ppo_update
from minijev.training.supervised import make_model, predict_logits

from .env import SnakeConfig, VectorSnakeEnv
from .evaluate import model_policy, play_games
from .task import SnakeTask

QUESTION = SnakeTask.default_question


def train_ppo(cfg: dict, checkpoint: str, run_dir: str | None = None, resume: bool = False) -> dict:
    rl = cfg.get("rl", {})
    seed = int(cfg.get("seed", 0))
    torch.manual_seed(seed)
    torch.set_num_threads(int(rl.get("threads", 8)))
    ck = torch.load(checkpoint, map_location="cpu", weights_only=False)
    base_cfg = ck["config"]
    task = SnakeTask(base_cfg.get("task", {}))  # the checkpoint decides the choice labels
    model = make_model(base_cfg, task)
    model.load_state_dict(ck["model"])
    model.log_temperature.zero_()
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model.to(device)
    ref = copy.deepcopy(model).eval()
    for p in ref.parameters():
        p.requires_grad_(False)

    run_dir = Path(run_dir or f"runs/{cfg.get('name', 'snake')}-ppo")
    log = RunLogger(run_dir, {**cfg, "init_checkpoint": checkpoint}, resume=resume)
    opt = torch.optim.AdamW(model.parameters(), lr=rl.get("lr", 1e-4), weight_decay=0.0)
    env_cfg = SnakeConfig.from_dict(cfg.get("env", {}))
    n_envs, T = int(rl.get("n_envs", 64)), int(rl.get("rollout", 128))
    total_updates = int(rl.get("total_env_steps", 1_000_000)) // (n_envs * T)
    start, best = 0, -1.0
    last_ck = run_dir / "checkpoints" / "rl_last.pt"
    if resume and last_ck.exists():
        st = torch.load(last_ck, weights_only=False)
        model.load_state_dict(st["model"])
        opt.load_state_dict(st["optimizer"])
        start, best = st["update"], st["best"]
        print(f"[rl] resumed at update {start}")
    envs = VectorSnakeEnv(n_envs, env_cfg, seed=seed * 7919 + start * n_envs)
    mcfg = base_cfg["model"]
    q = torch.as_tensor(encode_batch([QUESTION], mcfg.get("max_question_len", 32)), device=device).repeat(n_envs, 1)
    c = torch.as_tensor(encode_batch(list(task.default_choices), mcfg.get("max_choice_len", 16)),
                        device=device)[None].repeat(n_envs, 1, 1)
    gamma, lam = rl.get("gamma", 0.99), rl.get("gae_lambda", 0.95)
    warm = int(rl.get("value_warmup_updates", 5))
    eval_every = int(rl.get("eval_every", 10))
    log.write_json("env.json", {"env": environment_info(), "updates": total_updates, "n_envs": n_envs, "rollout": T})
    ep_scores: list[int] = []
    t0 = time.time()
    print(f"[rl] PPO: {total_updates} updates x {n_envs * T} steps, reward={cfg.get('env', {}).get('reward')}")
    for upd in range(start, total_updates):
        obs, acts, logps, vals, rews, dns = [], [], [], [], [], []
        model.eval()
        for _ in range(T):
            x = torch.as_tensor(task.encode_states(envs.states())["state_grid"], device=device)
            with torch.no_grad():
                out = model({"state_grid": x}, q, c)
            dist = torch.distributions.Categorical(logits=out["logits"])
            a = dist.sample()
            r, d, infos = envs.step(a.cpu().tolist())
            ep_scores += [i["score"] for i in infos if i]
            obs.append(x)
            acts.append(a)
            logps.append(dist.log_prob(a))
            vals.append(out["value"].cpu())
            rews.append(r)
            dns.append(d.astype(np.float32))
        with torch.no_grad():
            x = torch.as_tensor(task.encode_states(envs.states())["state_grid"], device=device)
            last_v = model({"state_grid": x}, q, c)["value"].cpu().numpy()
        V = torch.stack(vals).numpy()
        adv, ret = gae(np.stack(rews), V, np.stack(dns), last_v, gamma, lam)
        adv_t = torch.as_tensor(adv.reshape(-1), device=device)
        data = {"state": {"state_grid": torch.cat(obs)}, "question_ids": q.repeat(T, 1),
                "choice_ids": c.repeat(T, 1, 1), "actions": torch.cat(acts), "old_logp": torch.cat(logps),
                "adv": (adv_t - adv_t.mean()) / (adv_t.std() + 1e-8), "ret": torch.as_tensor(ret.reshape(-1), device=device)}
        stats = ppo_update(model, opt, data, rl, ref_model=ref, policy_coef=0.0 if upd < warm else 1.0)
        recent = ep_scores[-200:]
        log.log(upd + 1, env_steps=(upd + 1) * n_envs * T, train_avg_score=float(np.mean(recent)) if recent else 0.0,
                episodes=len(ep_scores), **{k: round(v, 5) for k, v in stats.items()})
        print(f"  update {upd + 1}/{total_updates} | sampled-policy avg score {np.mean(recent) if recent else 0:.1f} "
              f"| entropy {stats['entropy']:.3f} kl_ref {stats['kl_ref']:.4f} vf {stats['vf']:.4f} "
              f"| {time.time() - t0:.0f}s", flush=True)
        if (upd + 1) % eval_every == 0 or upd + 1 == total_updates:
            dec = TorchDecider(model=model, cfg=base_cfg, device=str(device))
            g = play_games(model_policy(dec), int(rl.get("eval_games", 100)), cfg.get("env", {}), seed=777)
            log.log(upd + 1, eval_avg_score=g["avg_score"], eval_fatal=g["fatal_move_rate"])
            print(f"  [eval] greedy avg score {g['avg_score']:.2f} median {g['median_score']} "
                  f"fatal-move {g['fatal_move_rate']:.4%}", flush=True)
            if g["avg_score"] > best:
                best = g["avg_score"]
                model.to(device)  # TorchDecider may have moved it
            torch.save({"model": model.state_dict(), "config": base_cfg, "update": upd + 1, "eval": g},
                           run_dir / "checkpoints" / "rl_best.pt")
            torch.save({"model": model.state_dict(), "optimizer": opt.state_dict(), "update": upd + 1,
                        "best": best}, last_ck)
    return finalize_rl(cfg, base_cfg, run_dir, log)


def finalize_rl(cfg, base_cfg, run_dir: Path, log) -> dict:
    """Re-calibrate the best RL policy on the validation split and evaluate it."""
    task = SnakeTask(base_cfg.get("task", {}))
    best = torch.load(run_dir / "checkpoints" / "rl_best.pt", weights_only=False)
    model = make_model(base_cfg, task)
    model.load_state_dict(best["model"])
    data_dir = Path(base_cfg["data"]["dir"])
    mc = base_cfg["model"]
    kw = dict(max_question_len=mc.get("max_question_len", 32), max_choice_len=mc.get("max_choice_len", 16))
    val = BatchEncoder(DecisionData.load(data_dir / "val.npz"), task, **kw)
    test = BatchEncoder(DecisionData.load(data_dir / "test.npz"), task, **kw)
    vl, vt = predict_logits(model, val, "cpu")
    tl, tt = predict_logits(model, test, "cpu")
    T = fit_temperature(vl, vt)
    model.log_temperature.fill_(float(np.log(T)))
    acc_ok = test.data.soft > 0 if test.data.soft is not None else None
    rep = calibration_report(tl, tt, T, acceptable=acc_ok)
    torch.save({"model": model.state_dict(), "config": base_cfg, "temperature": T}, run_dir / "checkpoints" / "final.pt")
    games = play_games(model_policy(TorchDecider(run_dir / "checkpoints" / "final.pt", threads=4)),
                       cfg.get("eval", {}).get("games", 200), cfg.get("env", {}))
    (run_dir / "games.json").write_text(json.dumps(games, indent=2))
    summary = {"params": model.num_parameters(), "temperature": T, "test_calibrated": rep,
               "test_raw": calibration_report(tl, tt, 1.0, acceptable=acc_ok), "games": games,
               "final_checkpoint": str(run_dir / "checkpoints" / "final.pt")}
    log.write_json("summary.json", summary)
    log.close()
    print(f"[rl] final: avg score {games['avg_score']:.2f} | oracle-agreement {rep['accuracy']:.4f} "
          f"| ECE vs oracle {rep['ece']:.4f} (T={T:.3f})")
    return summary


__all__ = ["train_ppo"]