#!/usr/bin/env python
"""Stage D: fit the value controller on logged counterfactual labels; calibrate on a held-out partition; tune
(lambda, kappa) on calibration ROLLOUTS (never on test); save the locked checkpoint.

Usage: python -m training.train_controller --config configs/experiments/controller_pilot.yaml [--device mps]
"""
from __future__ import annotations

import argparse, json, math, sys, time
from pathlib import Path

import numpy as np
import torch
import torch.nn.functional as F

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from prisme.losses import controller_losses, token_nll  # noqa: E402
from prisme.flops import action_cost_proxy  # noqa: E402
from training.common import Run, load_checkpoint, load_yaml, make_adamw, pick_device, save_checkpoint, seed_all, sync  # noqa: E402
from training.train_editor import build_model  # noqa: E402


def regenerate_state(model, cfg, rec, device):
    x = rec["tokens"].to(device).long()
    ctx = model.encode(x)
    S = ctx.S0
    B, T = x.shape
    r_pre, R_total = rec["r_pre"], rec["R_total"]
    for r in range(r_pre):
        a = rec["prefix"][:, :, r].to(device).long()
        S, _ = model.apply_action(S, ctx, a, r, R_total - r - 1)
    step = torch.full((B, T), r_pre, device=device, dtype=torch.long)
    remaining = torch.full((B, T), R_total - r_pre - 1, device=device, dtype=torch.long)
    return ctx, S, step, remaining


def preds_for_candidates(model, ctx, S, step, remaining, cand):
    p = model.controller(ctx.ctrl_in, S, step, remaining, ctx.entropy)              # [B,T,A]
    return {k: v.gather(-1, cand) for k, v in p.items()}, p


@torch.no_grad()
def evaluate_labels(model, cfg, recs, device):
    """Regression/classification diagnostics + oracle regret on audit batches."""
    model.eval()
    hub, sgn, brier, n = 0.0, 0.0, 0.0, 0
    regret, opp, top1, n_aud = 0.0, 0.0, 0.0, 0
    ece_bins = torch.zeros(10); ece_cnt = torch.zeros(10); ece_conf = torch.zeros(10)
    for rec in recs:
        ctx, S, step, remaining = regenerate_state(model, cfg, rec, device)
        cand = rec["cand"].to(device).long(); g = rec["gain"].to(device); v = rec["valid"].to(device)
        pc, pall = preds_for_candidates(model, ctx, S, step, remaining, cand)
        vv = v.unsqueeze(-1).expand_as(g)
        hub += float((F.huber_loss(pc["g_hat"], g, delta=0.1, reduction="none") * vv).sum()); n += int(vv.sum())
        sgn += float((((pc["g_hat"] > 0) == (g > 0)).float() * vv).sum())
        ph = torch.sigmoid(pc["harm_logit"]); harm = (g < 0).float()
        brier += float((((ph - harm) ** 2) * vv).sum())
        b = (ph * 10).clamp(max=9.999).long()[vv]
        ece_bins.index_add_(0, b.cpu(), harm[vv].cpu()); ece_conf.index_add_(0, b.cpu(), ph[vv].cpu()); ece_cnt.index_add_(0, b.cpu(), torch.ones_like(b, dtype=torch.float).cpu())
        if "all_gains" in rec:
            ag = rec["all_gains"].to(device)                      # [B,T,A]
            best = ag.max(-1).values.clamp(min=0)                 # one-step oracle value (STOP if nothing helps)
            q = pall["g_hat"]
            chosen = q.argmax(-1); chosen_q = q.max(-1).values
            realised = torch.where(chosen_q > 0, ag.gather(-1, chosen.unsqueeze(-1)).squeeze(-1), torch.zeros_like(best))
            regret += float(((best - realised) * v).sum()); opp += float((best * v).sum())
            top1 += float(((chosen == ag.argmax(-1)).float() * v).sum()); n_aud += int(v.sum())
    ece = float((ece_cnt / ece_cnt.sum().clamp(min=1) * (ece_bins / ece_cnt.clamp(min=1) - ece_conf / ece_cnt.clamp(min=1)).abs()).sum())
    out = {"huber": hub / max(n, 1), "sign_acc": sgn / max(n, 1), "harm_brier": brier / max(n, 1), "harm_ece": ece, "n_labels": n}
    if n_aud:
        out.update({"oracle_opportunity_per_state": opp / n_aud, "regret_per_state": regret / n_aud,
                    "captured_fraction": 1 - regret / max(opp, 1e-9), "top1_match": top1 / n_aud, "n_audit_states": n_aud})
    return out


@torch.no_grad()
def rollout_cost_quality(model, cfg, recs, device, cost, lam, kap, max_budget):
    """Deployed-controller rollout (teacher forced) on calibration batches: mean loss, mean edits, mean ops ratio."""
    tot, n, edits, ops = 0.0, 0, 0.0, 0.0
    base_ops = 1.0
    for rec in recs:
        x = rec["tokens"].to(device).long(); y = rec["targets"].to(device).long()
        out = model.forward_policy(x, max_budget, cost, lambda_cost=lam, kappa_risk=kap)
        nll = token_nll(out["logits"], y); v = (y != -100)
        tot += float(nll[v].sum()); n += int(v.sum()); edits += float(out["n_edits"][v].float().sum())
        if out["actions"] is not None:
            a = out["actions"][v]                                     # [n, R]
            ops += float(torch.where(a < cfg.n_actions, cost[a.clamp(max=cfg.n_actions - 1)], torch.zeros_like(cost[0]).expand_as(a).float()).sum())
    return {"loss": tot / max(n, 1), "edits_per_token": edits / max(n, 1), "extra_ops_ratio": ops / max(n, 1)}


def main(argv=None) -> int:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--config", required=True)
    ap.add_argument("--device", default="auto")
    ap.add_argument("--seed", type=int, default=None)
    ap.add_argument("--run-root", default="experiments/runs")
    args = ap.parse_args(argv)
    exp = load_yaml(args.config)
    seed = args.seed if args.seed is not None else exp.get("seed", 1)
    seed_all(seed); device = pick_device(args.device)
    cfg, model = build_model(exp)
    model = model.to(device)
    meta = load_checkpoint(exp["checkpoint"], model, strict=False)
    for n_, p in model.named_parameters():
        p.requires_grad_(n_.startswith("controller"))
    train = torch.load(exp["labels"]["train"], weights_only=False)
    calib = torch.load(exp["labels"]["calib"], weights_only=False)
    assert train["checkpoint_sha256"] == calib["checkpoint_sha256"] == meta.get("weights_sha256"), "labels were generated by a different editor checkpoint (stale labels)"
    tr = exp["train"]
    run = Run(exp.get("name", "controller"), {"experiment": exp, "model": cfg.to_dict()}, root=args.run_root, seed=seed)
    run.write_manifest(device=str(device), checkpoint_parent=exp["checkpoint"], labels_train=exp["labels"]["train"], labels_calib=exp["labels"]["calib"])
    opt = make_adamw(model.controller.parameters(), tr["lr"], tr.get("weight_decay", 0.01))
    recs = train["records"]; rng = np.random.default_rng(seed)
    alpha, beta, gamma = tr.get("alpha", 1.0), tr.get("beta", 0.5), tr.get("gamma", 0.1)
    model.train()
    best_state, best_score, best_step = None, -1e9, -1
    def score_of(ev):
        return -ev["regret_per_state"] if "regret_per_state" in ev else ev["sign_acc"]
    for step in range(tr["steps"]):
        rec = recs[int(rng.integers(0, len(recs)))]
        with torch.no_grad():
            ctx, S, st, rem = regenerate_state(model, cfg, rec, device)
        cand = rec["cand"].to(device).long(); g = rec["gain"].to(device); v = rec["valid"].to(device).unsqueeze(-1).expand_as(g)
        pc, pall = preds_for_candidates(model, ctx, S, st, rem, cand)
        L = controller_losses({k: x_[v] for k, x_ in pc.items()}, g[v], torch.ones_like(g[v], dtype=torch.bool), tr.get("huber_delta", 0.1))
        loss = alpha * L["gain"] + beta * L["harm"] + gamma * L["nll"]
        # ---- within-state ranking losses (D12): pairwise on the 4 logged candidates, listwise on exhaustive audits
        w_pair, w_list = tr.get("pairwise", 0.0), tr.get("listwise", 0.0)
        if w_pair > 0:
            gi, gj = g.unsqueeze(-1), g.unsqueeze(-2)                       # [B,T,4,4]
            pi, pj = pc["g_hat"].unsqueeze(-1), pc["g_hat"].unsqueeze(-2)
            wgt = (gi - gj).abs() * rec["valid"].to(device).float()[..., None, None]
            lp = F.binary_cross_entropy_with_logits((pi - pj) / tr.get("rank_temp", 0.02), (gi > gj).float(), weight=wgt, reduction="sum") / wgt.sum().clamp(min=1e-6)
            loss = loss + w_pair * lp; L["pair"] = lp
        if w_list > 0 and "all_gains" in rec:
            ag = rec["all_gains"].to(device)                                  # [B,T,A]
            tgt = torch.softmax(ag / tr.get("rank_temp", 0.02), -1)
            ll = -(tgt * torch.log_softmax(pall["g_hat"] / tr.get("rank_temp", 0.02), -1)).sum(-1)
            ll = (ll * rec["valid"].to(device).float()).sum() / rec["valid"].sum().clamp(min=1)
            loss = loss + w_list * ll; L["list"] = ll
        opt.zero_grad(set_to_none=True); loss.backward()
        torch.nn.utils.clip_grad_norm_(model.controller.parameters(), 1.0); opt.step()
        if step % tr.get("log_every", 50) == 0 or step == tr["steps"] - 1:
            run.log(step=step, loss=float(loss), gain=float(L["gain"]), harm=float(L["harm"]), nll=float(L["nll"]))
            print(f"step {step} loss {float(loss):.5f} gain {float(L['gain']):.5f} harm {float(L['harm']):.4f}", flush=True)
        if (step + 1) % tr.get("eval_every", 500) == 0 or step == tr["steps"] - 1:
            ev = evaluate_labels(model, cfg, calib["records"], device); model.train()
            run.log(step=step, **{f"calib_{k}": v_ for k, v_ in ev.items()}); print("  calib:", {k: round(v_, 4) for k, v_ in ev.items()}, flush=True)
            if score_of(ev) > best_score:
                best_score, best_step = score_of(ev), step + 1
                best_state = {k: v_.detach().cpu().clone() for k, v_ in model.controller.state_dict().items()}
    if best_state is not None:   # early stopping on the calibration partition (never on test)
        model.controller.load_state_dict(best_state)
        run.log(event="restored_best_calibration_checkpoint", step=best_step, score=best_score)
        print(f"restored controller from step {best_step} (calibration score {best_score:.5f})", flush=True)
    model.eval()
    ev_train = evaluate_labels(model, cfg, recs[: min(len(recs), 40)], device)
    ev_calib = evaluate_labels(model, cfg, calib["records"], device)
    # ---- (lambda, kappa) sweep by calibration rollouts ------------------------------------------------
    cost = torch.tensor(action_cost_proxy(cfg, tr["seq_len"]), device=device)
    sweep = []
    cal_sub = calib["records"][: exp.get("sweep_batches", 10)]
    for lam in exp.get("lambda_grid", [0.0, 0.01, 0.03, 0.1]):
        for kap in exp.get("kappa_grid", [0.0, 0.5, 1.0]):
            r = rollout_cost_quality(model, cfg, cal_sub, device, cost, lam, kap, exp.get("max_budget", 8))
            sweep.append({"lambda": lam, "kappa": kap, **r}); print("sweep", sweep[-1], flush=True)
    b0 = rollout_cost_quality(model, cfg, cal_sub, device, cost, 0.0, 0.0, 0)
    sel_rule = exp.get("selection_rule", "max_loss_reduction_per_ops")
    for s in sweep:
        s["loss_reduction"] = b0["loss"] - s["loss"]
        s["reduction_per_ops"] = s["loss_reduction"] / max(s["extra_ops_ratio"], 1e-9)
    if sel_rule == "max_loss_reduction_per_ops":
        cands = [s for s in sweep if s["loss_reduction"] > 0] or sweep
        best = max(cands, key=lambda s: s["reduction_per_ops"])
    else:
        best = min(sweep, key=lambda s: s["loss"])
    cfg.lambda_cost, cfg.kappa_risk = best["lambda"], best["kappa"]
    sha = save_checkpoint(run.dir / "checkpoint", model, cfg.to_dict(), {"run_id": run.run_id, "parent": exp["checkpoint"], "selected": best, "locked": True})
    run.finish(eval_train=ev_train, eval_calib=ev_calib, sweep=sweep, baseline_budget0=b0, selected=best, checkpoint_sha256=sha, selection_rule=sel_rule)
    print("selected", best, "\nrun", run.run_id)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())