"""Head-RL PPO — train PolicyHeadV2 with DAUS-reward PPO on tick-level rollouts.

Episode = one video clip (group of ticks sharing the same video_id).
Step    = one tick within the clip.
State   = (BELIEF features, danger features, prev_action_embed) for that tick.
Action  ∈ {SILENT=0, OBSERVE=1, ALERT=2}
Reward  = DAUS-B' contribution per-clip, distributed equally across ticks.

PPO loss:
  L_clip = E[min(r_t · A_t, clip(r_t, 1-ε, 1+ε) · A_t)]
  where r_t = π_θ(a|s) / π_θ_old(a|s),  A_t = reward − value(s)
  + value-loss + entropy bonus + KL penalty to ref policy

Since our cache is tick-level (not per-clip-trajectory), we approximate the
episode by grouping ticks by video_id. For simplicity (and because the user's
9k legacy is mostly 1 tick per video), we treat each tick as a 1-step episode
with reward defined by the DAUS-style preference (analogous to bandit RL).

Reward shaping (per tick, based on pair_type for this tick if available):
  - tta < 0.5s & action=ALERT     → +1.0 (correct ALERT)
  - tta < 0.5s & action=OBSERVE   → +0.3
  - tta ∈ [4, 6]s & action=OBSERVE→ +0.8 (correct borderline OBSERVE)
  - tta ∈ [4, 6]s & action=SILENT → -0.3 (missed borderline)
  - tta ∈ [1.5, 2]s & action=OBSERVE→ +0.5 (correct de-escalation)
  - tta ∈ [1.5, 2]s & action=ALERT→ -0.2 (over-confident)
  - tta > 8s & action=SILENT      → +0.5 (correct safe)
  - any false ALERT on safe       → -0.5
  + KL penalty: -β·KL(π_θ || π_ref)

Usage:
    python -m training.Policy.train_head_ppo \
        --train_cache data/belief_cache_v3/sft_x_v3__train_9k.pt \
        --policy_warm checkpoints/policy_v3_strong/best.pt \
        --out_dir checkpoints/policy_v3_head_ppo
"""
from __future__ import annotations

import argparse
import json
import logging
import sys
from pathlib import Path

import torch
import torch.nn.functional as F
from torch.utils.data import DataLoader, Dataset
from tqdm import tqdm

ROOT = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(ROOT))
from lkalert.models.danger_head import DangerHead
from lkalert.models.policy_head_v2 import PolicyHeadV2
from training.Policy._balance_eval import evaluate_policy_on_val, format_gate_row

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
logger = logging.getLogger("head_ppo")


def reward_fn(tick_action: int, action: int, tta_raw: float) -> float:
    """Bandit reward for a (tick, action) pair given the ground-truth context.

    Returns a scalar reward in [-1, +1].
    """
    is_alert_gt = (tick_action == 2)
    is_obs_gt   = (tick_action == 1)
    is_sil_gt   = (tick_action == 0)

    # Hard-coded reward shaping per pair-type criteria
    if tta_raw is None or tta_raw < 0:
        # safe_neg (no event)
        if action == 0: return +0.5
        if action == 1: return -0.2  # false OBSERVE on truly safe
        if action == 2: return -0.5  # false ALERT
        return 0.0

    if tta_raw <= 0.5:        # imminent collision
        if action == 2: return +1.0
        if action == 1: return +0.3
        return -0.8           # missed imminent ALERT
    if tta_raw < 2.0:         # near event
        if action == 2: return +0.7
        if action == 1: return +0.5
        return -0.5
    if tta_raw < 4.0:         # OBSERVE window
        if action == 1: return +0.8
        if action == 2: return +0.0
        return -0.3
    if tta_raw < 6.0:         # borderline OBSERVE/SILENT
        if action == 1: return +0.5
        if action == 0: return -0.1
        return -0.4           # premature ALERT
    if tta_raw >= 8.0:        # clearly far
        if action == 0: return +0.5
        if action == 1: return -0.2
        return -0.5
    return 0.0


class CacheDataset(Dataset):
    """Iterate all cache ticks (skip ones with no GT tta)."""
    def __init__(self, cache_path: Path, observe_oversample: int = 4):
        self.cache = torch.load(cache_path, weights_only=False, map_location="cpu")
        # Build index list with optional OBSERVE oversample
        ta = self.cache["tick_action"]
        idxs = []
        for i in range(len(ta)):
            idxs.append(i)
            if ta[i] == 1:
                idxs.extend([i] * (observe_oversample - 1))
        self.indices = idxs
        logger.info(f"  cache N={len(ta)}  oversampled OBSERVE × {observe_oversample} → total {len(idxs)}")

    def __len__(self): return len(self.indices)
    def __getitem__(self, idx):
        ci = self.indices[idx]
        return {
            "belief": self.cache["belief_content"][ci],
            "policy": self.cache["policy_position"][ci],
            "valid": self.cache["valid_frames"][ci],
            "tick_action": int(self.cache["tick_action"][ci]),
            "tta_raw": float(self.cache.get("tick_tta_raw",
                              torch.full((len(self.cache["tick_action"]),), -1.0))[ci]),
        }


def collate(batch):
    return {
        "belief": torch.stack([b["belief"] for b in batch]),
        "policy": torch.stack([b["policy"] for b in batch]),
        "valid": torch.stack([b["valid"] for b in batch]),
        "tick_action": torch.tensor([b["tick_action"] for b in batch], dtype=torch.long),
        "tta_raw": torch.tensor([b["tta_raw"] for b in batch], dtype=torch.float32),
    }


def ppo_step(policy, ref_policy, batch, opt, dh, device,
              clip_ratio=0.2, entropy_coef=0.01, kl_coef=0.05,
              rl_weight=0.7, alert_anchor_weight=1.0, reward_bonus_alert=0.0):
    bc = batch["belief"].to(device, dtype=torch.float32, non_blocking=True)
    pp = batch["policy"].to(device, dtype=torch.float32, non_blocking=True)
    v  = batch["valid"].to(device, non_blocking=True)
    ta = batch["tick_action"].to(device, non_blocking=True)
    tta = batch["tta_raw"].to(device, non_blocking=True)
    B = bc.shape[0]
    prev = torch.full((B,), 3, dtype=torch.long, device=device)

    with torch.no_grad():
        dh_out = dh(bc, valid_frames=v)
        ref_logits = ref_policy(pp, dh_out["perception_summary"],
                                  dh_out["per_frame"], prev, valid_frames=v)
        ref_log_p = F.log_softmax(ref_logits, dim=-1)
        ref_p = ref_log_p.exp()
        old_log_p = ref_log_p.detach()
        actions = torch.multinomial(ref_p, 1).squeeze(-1)
        old_logp_a = ref_log_p.gather(1, actions.unsqueeze(1)).squeeze(1)

    # Reward via reward_fn + bonus for correct ALERT on real-ALERT tick
    rewards_list = []
    for i in range(B):
        r = reward_fn(int(ta[i].item()), int(actions[i].item()),
                      tta[i].item() if tta[i].item() >= 0 else None)
        # Extra bonus for correct ALERT predictions on real ALERT
        if reward_bonus_alert > 0 and int(ta[i].item()) == 2 and int(actions[i].item()) == 2:
            r += reward_bonus_alert
        rewards_list.append(r)
    rewards = torch.tensor(rewards_list, dtype=torch.float32, device=device)

    adv = rewards - rewards.mean()
    if adv.std() > 1e-6:
        adv = adv / (adv.std() + 1e-6)

    logits = policy(pp, dh_out["perception_summary"], dh_out["per_frame"],
                    prev, valid_frames=v)
    log_p = F.log_softmax(logits, dim=-1)
    new_logp_a = log_p.gather(1, actions.unsqueeze(1)).squeeze(1)
    ratio = (new_logp_a - old_logp_a).exp()
    s1 = ratio * adv
    s2 = ratio.clamp(1 - clip_ratio, 1 + clip_ratio) * adv
    ppo_loss = -torch.min(s1, s2).mean()

    entropy = -(log_p.exp() * log_p).sum(dim=-1).mean()
    kl = (log_p.exp() * (log_p - ref_log_p)).sum(dim=-1).mean()

    rl_term = ppo_loss - entropy_coef * entropy + kl_coef * kl

    # ALERT anchor CE on real-ALERT samples
    anchor_mask = (ta == 2)
    if anchor_mask.any():
        anchor_l = F.cross_entropy(logits[anchor_mask], ta[anchor_mask])
    else:
        anchor_l = torch.zeros((), device=device)

    total = (rl_weight * rl_term
             + (1 - rl_weight) * alert_anchor_weight * anchor_l)
    total.backward()
    torch.nn.utils.clip_grad_norm_(policy.parameters(), 1.0)
    opt.step(); opt.zero_grad(set_to_none=True)

    with torch.no_grad():
        mean_reward = rewards.mean().item()
        return {
            "loss": total.item(), "ppo": ppo_loss.item(),
            "anchor": anchor_l.item(),
            "entropy": entropy.item(), "kl": kl.item(),
            "reward": mean_reward,
        }


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--train_cache", type=Path,
                    default=ROOT / "data/belief_cache_v3/sft_x_v3__train_9k.pt")
    ap.add_argument("--val_cache", type=Path,
                    default=ROOT / "data/belief_cache_v3/sft_x_v3__multisrc_val.pt")
    ap.add_argument("--policy_warm", type=Path,
                    default=ROOT / "checkpoints/policy_v3_strong/best.pt")
    ap.add_argument("--danger_ckpt", type=Path,
                    default=ROOT / "checkpoints/danger_v2/seed2/best.pt")
    ap.add_argument("--out_dir", type=Path, required=True)
    ap.add_argument("--lr", type=float, default=1e-5)
    ap.add_argument("--epochs", type=int, default=10)
    ap.add_argument("--batch_size", type=int, default=64)
    ap.add_argument("--clip_ratio", type=float, default=0.2)
    ap.add_argument("--kl_coef", type=float, default=0.1)
    ap.add_argument("--entropy_coef", type=float, default=0.01)
    ap.add_argument("--observe_oversample", type=int, default=3)
    ap.add_argument("--rl_weight", type=float, default=0.7)
    ap.add_argument("--alert_anchor_weight", type=float, default=1.0)
    ap.add_argument("--reward_bonus_alert", type=float, default=2.0,
                    help="Extra reward on correct ALERT prediction at real ALERT tick")
    ap.add_argument("--max_samples", type=int, default=0)
    ap.add_argument("--seed", type=int, default=0)
    args = ap.parse_args()
    args.out_dir.mkdir(parents=True, exist_ok=True)

    torch.manual_seed(args.seed)
    device = "cuda" if torch.cuda.is_available() else "cpu"

    ds = CacheDataset(args.train_cache, observe_oversample=args.observe_oversample)
    if args.max_samples > 0:
        ds.indices = ds.indices[:args.max_samples]
        logger.info(f"  truncated to {len(ds.indices)} samples (smoke)")
    loader = DataLoader(ds, batch_size=args.batch_size, shuffle=True,
                        num_workers=2, collate_fn=collate, pin_memory=True)

    ck_d = torch.load(args.danger_ckpt, weights_only=False, map_location="cpu")
    dh = DangerHead(in_dim=ck_d["in_dim"]).to(device)
    dh.load_state_dict(ck_d["model"]); dh.eval()
    for p in dh.parameters(): p.requires_grad_(False)

    ck_p = torch.load(args.policy_warm, weights_only=False, map_location="cpu")
    ph_kwargs = dict(
        policy_dim=ck_p.get("policy_dim", 2560),
        perception_dim_per_query=ck_p.get("perception_dim_per_query", 512),
        k_queries=ck_p.get("k_queries", 4),
    )
    policy = PolicyHeadV2(**ph_kwargs).to(device)
    policy.load_state_dict(ck_p["model"])

    ref_policy = PolicyHeadV2(**ph_kwargs).to(device)
    ref_policy.load_state_dict(ck_p["model"])
    ref_policy.eval()
    for p in ref_policy.parameters(): p.requires_grad_(False)

    val_cache = None
    if args.val_cache.exists() and args.epochs >= 1 and args.max_samples == 0:
        logger.info(f"[load] val_cache={args.val_cache}")
        val_cache = torch.load(args.val_cache, weights_only=False, map_location="cpu")
        logger.info(f"  val N={len(val_cache['ids'])}")

    opt = torch.optim.AdamW(policy.parameters(), lr=args.lr, weight_decay=1e-5)

    best_composite = -1e9
    log_records = []
    for ep in range(args.epochs):
        policy.train()
        run = {"loss": 0, "ppo": 0, "anchor": 0, "entropy": 0, "kl": 0, "reward": 0}
        n_b = 0
        pbar = tqdm(loader, ncols=80, desc=f"ppo ep{ep}")
        for b in pbar:
            stats = ppo_step(policy, ref_policy, b, opt, dh, device,
                              clip_ratio=args.clip_ratio,
                              entropy_coef=args.entropy_coef,
                              kl_coef=args.kl_coef,
                              rl_weight=args.rl_weight,
                              alert_anchor_weight=args.alert_anchor_weight,
                              reward_bonus_alert=args.reward_bonus_alert)
            for k, v in stats.items():
                run[k] += v
            n_b += 1
            pbar.set_postfix(reward=run["reward"]/n_b, kl=run["kl"]/n_b,
                             anc=run["anchor"]/n_b)

        rec = {"epoch": ep, **{k: v / max(1, n_b) for k, v in run.items()}}

        if val_cache is not None:
            val_m = evaluate_policy_on_val(policy, dh, val_cache, device,
                                            batch_size=256)
            rec["val"] = val_m
            logger.info(format_gate_row(val_m, tag=f"ppo ep{ep}"))
            composite = val_m["composite"]
            if composite > best_composite:
                best_composite = composite
                save_dict = {
                    "model": policy.state_dict(),
                    "policy_dim": ph_kwargs["policy_dim"],
                    "perception_dim_per_query": ph_kwargs["perception_dim_per_query"],
                    "k_queries": ph_kwargs["k_queries"],
                    "args": vars(args), "epoch": ep,
                    "val_metrics": val_m, "composite": composite,
                }
                torch.save(save_dict, args.out_dir / "best.pt")
                logger.info(f"  [save best] composite={composite:.4f}")
        else:
            save_dict = {
                "model": policy.state_dict(),
                "policy_dim": ph_kwargs["policy_dim"],
                "perception_dim_per_query": ph_kwargs["perception_dim_per_query"],
                "k_queries": ph_kwargs["k_queries"],
                "args": vars(args), "epoch": ep,
            }
            torch.save(save_dict, args.out_dir / "best.pt")

        log_records.append(rec)
        logger.info(f"[ep{ep}] loss={rec['loss']:.4f} reward={rec['reward']:.4f} "
                    f"anchor={rec['anchor']:.4f} kl={rec['kl']:.4f}")

    (args.out_dir / "training_log.json").write_text(
        json.dumps(log_records, indent=2, default=str))
    logger.info(f"\n[done] best composite={best_composite:.4f}  saved to {args.out_dir}/best.pt")


if __name__ == "__main__":
    main()