"""Task-agnostic PPO update for MiniJEV (actor = choice logits, critic = CLS value head).

Why PPO (and not DQN) for MiniJEV:
  * the model already outputs a distribution over *dynamic* choices; PPO trains
    exactly that distribution, DQN would turn scores into Q-values and lose the
    probabilistic reading that calibration relies on;
  * it starts cleanly from the supervised policy (same head), with a KL anchor to
    the supervised reference so early noisy advantages do not wreck it;
  * on-policy + vectorised envs keeps it simple and stable at small scale.
"""
from __future__ import annotations

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


def gae(rewards: np.ndarray, values: np.ndarray, dones: np.ndarray, last_value: np.ndarray,
        gamma: float, lam: float) -> tuple[np.ndarray, np.ndarray]:
    """rewards/values/dones: (T, N). Returns (advantages, returns)."""
    T = rewards.shape[0]
    adv = np.zeros_like(rewards)
    last = np.zeros_like(last_value)
    for t in reversed(range(T)):
        nv = last_value if t == T - 1 else values[t + 1]
        nonterminal = 1.0 - dones[t]
        delta = rewards[t] + gamma * nv * nonterminal - values[t]
        last = delta + gamma * lam * nonterminal * last
        adv[t] = last
    return adv, adv + values


def ppo_update(model, opt, data: dict, cfg: dict, ref_model=None, policy_coef: float = 1.0) -> dict:
    """data: state (dict of tensors), question_ids, choice_ids, actions, old_logp, adv, ret (flattened)."""
    n = len(data["actions"])
    mb = int(cfg.get("minibatch", 1024))
    clip = cfg.get("clip", 0.2)
    stats = {"pg": 0.0, "vf": 0.0, "entropy": 0.0, "kl_ref": 0.0, "approx_kl": 0.0, "clipfrac": 0.0}
    count = 0
    model.train()
    for _ in range(int(cfg.get("epochs", 2))):
        perm = torch.randperm(n, device=data["actions"].device)
        for i in range(0, n, mb):
            idx = perm[i:i + mb]
            state = {k: v[idx] for k, v in data["state"].items()}
            out = model(state, data["question_ids"][idx], data["choice_ids"][idx])
            logp_all = F.log_softmax(out["logits"], -1)
            logp = logp_all.gather(1, data["actions"][idx, None]).squeeze(1)
            ratio = (logp - data["old_logp"][idx]).exp()
            adv = data["adv"][idx]
            pg = -torch.min(ratio * adv, ratio.clamp(1 - clip, 1 + clip) * adv).mean()
            vf = 0.5 * (out["value"] - data["ret"][idx]).pow(2).mean()
            p = logp_all.exp()
            ent = -(p * logp_all.masked_fill(~out["choice_mask"], 0.0)).sum(-1).mean()
            loss = policy_coef * pg + cfg.get("vf_coef", 0.5) * vf - cfg.get("ent_coef", 0.001) * ent
            kl_ref = torch.zeros(())
            if ref_model is not None and cfg.get("kl_ref_coef", 0.0) > 0:
                with torch.no_grad():
                    ref_logp = F.log_softmax(ref_model(state, data["question_ids"][idx],
                                                       data["choice_ids"][idx])["logits"], -1)
                kl_ref = (ref_logp.exp() * (ref_logp - logp_all)).masked_fill(~out["choice_mask"], 0.0).sum(-1).mean()
                loss = loss + cfg["kl_ref_coef"] * kl_ref
            opt.zero_grad(set_to_none=True)
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), cfg.get("grad_clip", 0.5))
            opt.step()
            with torch.no_grad():
                stats["pg"] += float(pg)
                stats["vf"] += float(vf)
                stats["entropy"] += float(ent)
                stats["kl_ref"] += float(kl_ref)
                stats["approx_kl"] += float((data["old_logp"][idx] - logp).mean())
                stats["clipfrac"] += float(((ratio - 1).abs() > clip).float().mean())
            count += 1
    return {k: v / max(1, count) for k, v in stats.items()}