#!/usr/bin/env python3
"""LKAlert-BD multi-horizon trainer (Day 3).

Extends `POMDPTemporalHead` with:
  * a dynamic-feature side-channel built by `dynamic_features.build_features`
  * configurable multi-horizon binary outputs `{p_any, p_1500, p_1000, p_500,
    p_resolution_proxy, p_ego}`
  * optional ordinal monotonic regulariser `p_500 ≤ p_1000 ≤ p_1500 ≤ p_any`
  * automatic skipping of heads whose training labels are degenerate
    (single class) — e.g. the current Nexar train cache has 0 clips with
    TTA ≤ 1.0 s, so p_1000 and p_500 are unsupported and silently dropped.

Inputs / labels:
  * belief cache: `data/belief_cache_perframe_qwen3vl4b/nexar_train_diag.pt`
    keys: beliefs_frame [N,T,D], valid_frames [N,T], beliefs_text [N,D],
          tta_means [N], tta_vars [N], meta.{ids, action_labels}
  * diag json:    `data/policy_labels/nexar_train_diag.json`
                  per-clip {tta_raw, action_label, category, source}

Reused frozen heads remain on disk; training is ~6 min × 5 seeds on CPU.

Output:
  checkpoints/Policy/lkalert_bd_seed{0..4}/best.pt
  checkpoints/Policy/lkalert_bd_best/  (symlink to best by nexar_val ap)
"""
from __future__ import annotations

import argparse
import json
import logging
import random
import sys
from pathlib import Path
from typing import Dict, List, Optional, Tuple

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

sys.path.insert(0, str(Path(__file__).resolve().parents[2]))

from training.Policy.dynamic_features import build_features, feature_dim
from training.Policy.train_pomdp_head import POMDPTemporalHead

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

CACHE_DIR = Path("data/belief_cache_perframe_qwen3vl4b")
DIAG_DIR = Path("data/policy_labels")


def _diag_filename(cache_name: str) -> str:
    """Map cache stem to its diag-json filename.

    Convention: most caches use `{cache}_diag.json` but `nexar_train_diag`
    already has `_diag` in its name, so the file is `nexar_train_diag.json`.
    """
    if cache_name.endswith("_diag"):
        return f"{cache_name}.json"
    return f"{cache_name}_diag.json"

HORIZON_BUCKETS = [
    ("p_500",  0.0, 0.5),
    ("p_1000", 0.0, 1.0),
    ("p_1500", 0.0, 1.5),
    ("p_any",  0.0, 2.5),
]


# ─── label derivation ────────────────────────────────────────────────────────

def derive_labels(diag_json: Path, ids: List[str]) -> Dict[str, np.ndarray]:
    """Return per-clip binary label arrays keyed by horizon name + p_ego +
    p_resolution_proxy. Aligned to the cache's id order.

    `p_any` is `action_label != 0` (consistent across Nexar / DoTA / DAD /
    DADA, where DoTA-style caches set tta_raw = -1 even for true positives).

    `p_1500` / `p_1000` / `p_500` derive from `tta_raw` and require
    `tta_raw >= 0`; otherwise the bucket label is 0 (caches without TTA
    info simply produce all-0 horizon labels and the trainer skips them).
    """
    raw = json.loads(diag_json.read_text())
    by_id = {s["video_id"]: s for s in raw["samples"]}
    tta = np.asarray([by_id[v]["tta_raw"]      for v in ids], dtype=np.float32)
    cat = [by_id[v]["category"]                for v in ids]
    al  = np.asarray([by_id[v]["action_label"] for v in ids], dtype=np.int32)

    labels: Dict[str, np.ndarray] = {}
    has_tta = tta >= 0.0
    for name, lo, hi in HORIZON_BUCKETS:
        if name == "p_any":
            # cross-domain compatible: any non-SILENT action
            labels[name] = (al != 0).astype(np.float32)
        else:
            # explicit TTA bucket; requires tta_raw meaningful
            labels[name] = (has_tta & (tta >= lo) & (tta <= hi)).astype(np.float32)
    labels["p_ego"] = np.asarray([1.0 if c == "ego_positive" else 0.0
                                   for c in cat], dtype=np.float32)
    labels["p_resolution_proxy"] = 1.0 - labels["p_any"]
    return labels


def usable_heads(labels: Dict[str, np.ndarray], head_names: List[str]
                  ) -> List[str]:
    out = []
    for name in head_names:
        y = labels[name]
        if 0 < y.sum() < y.size:
            out.append(name)
        else:
            logger.warning(f"  skip {name}: degenerate (n_pos={int(y.sum())} "
                           f"of {y.size})")
    return out


# ─── dataset ────────────────────────────────────────────────────────────────

class CacheDataset(Dataset):
    def __init__(self, cache_path: Path, label_path: Optional[Path],
                 head_names: List[str]):
        d = torch.load(cache_path, weights_only=False, map_location="cpu")
        self.bf = d["beliefs_frame"].float()
        self.vf = d["valid_frames"].bool()
        self.bt = d["beliefs_text"].float()
        self.tm = d["tta_means"].float()
        self.tv = d["tta_vars"].float()
        self.ids = d["meta"]["ids"]
        if label_path is not None:
            self.lbl = derive_labels(label_path, self.ids)
        else:
            # eval cache without diag: use cache.action_labels for p_any only
            al = np.asarray(d["meta"].get("action_labels", []),
                             dtype=np.int32)
            self.lbl = {"p_any": (al == 2).astype(np.float32)}
        self.head_names = head_names

    def __len__(self) -> int:
        return self.bf.shape[0]

    def __getitem__(self, i: int) -> Dict[str, torch.Tensor]:
        out = {
            "belief": self.bf[i],  "valid": self.vf[i],
            "text":   self.bt[i],  "tta_mean": self.tm[i:i+1].squeeze(0),
            "tta_var": self.tv[i:i+1].squeeze(0),
            "vid":    self.ids[i],
        }
        for h in self.head_names:
            if h in self.lbl:
                out[f"y_{h}"] = torch.tensor(self.lbl[h][i], dtype=torch.float32)
        return out


def collate(batch: List[Dict]) -> Dict[str, torch.Tensor]:
    out = {
        "belief":   torch.stack([b["belief"]   for b in batch]),
        "valid":    torch.stack([b["valid"]    for b in batch]),
        "text":     torch.stack([b["text"]     for b in batch]),
        "tta_mean": torch.stack([b["tta_mean"] for b in batch]),
        "tta_var":  torch.stack([b["tta_var"]  for b in batch]),
        "vids":     [b["vid"] for b in batch],
    }
    for k in batch[0]:
        if k.startswith("y_"):
            out[k] = torch.stack([b[k] for b in batch])
    return out


# ─── model ──────────────────────────────────────────────────────────────────

class LKAlertBDHead(nn.Module):
    """POMDP trunk + dynamic-feature side-channel + per-horizon binary heads."""

    def __init__(self, in_dim: int = 2560, proj_dim: int = 512,
                 gru_hidden: int = 256, dyn_dim: int = 10250,
                 dyn_hidden: int = 64, dropout: float = 0.2,
                 head_names: Optional[List[str]] = None):
        super().__init__()
        self.head_names = head_names or ["p_any"]

        # belief trunk (mirrors POMDPTemporalHead exactly so we can warm-start)
        self.in_proj = nn.Sequential(
            nn.Linear(in_dim, proj_dim),
            nn.LayerNorm(proj_dim),
            nn.GELU(),
            nn.Dropout(dropout),
        )
        self.text_proj = nn.Sequential(
            nn.Linear(in_dim, gru_hidden),
            nn.LayerNorm(gru_hidden),
            nn.Tanh(),
        )
        self.gru = nn.GRU(proj_dim, gru_hidden, num_layers=1, batch_first=True)
        self.attn = nn.Linear(gru_hidden, 1)

        # dynamic-feature side-channel
        self.dyn_proj = nn.Sequential(
            nn.Linear(dyn_dim, 256),
            nn.LayerNorm(256),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(256, dyn_hidden),
            nn.GELU(),
        )
        # Larger BD trunk to give the side-channel a meaningful effect.
        self.fusion = nn.Sequential(
            nn.Linear(gru_hidden + dyn_hidden, 128),
            nn.GELU(),
            nn.Dropout(dropout),
        )
        # one binary head per supported horizon
        self.heads = nn.ModuleDict({
            h: nn.Linear(128, 1) for h in self.head_names
        })

    def trunk_pool(self, beliefs: torch.Tensor, valid: torch.Tensor,
                    text: torch.Tensor) -> torch.Tensor:
        x = self.in_proj(beliefs)
        h0 = self.text_proj(text).unsqueeze(0).contiguous()
        out, _ = self.gru(x, h0)
        attn_logits = self.attn(out).squeeze(-1)
        attn_logits = attn_logits.masked_fill(~valid, float("-inf"))
        empty = (~valid).all(dim=1)
        if empty.any():
            attn_logits[empty] = 0.0
        w = F.softmax(attn_logits, dim=1).unsqueeze(-1)
        return (out * w).sum(dim=1)                                # [B, H]

    def forward(self, beliefs: torch.Tensor, valid: torch.Tensor,
                text: torch.Tensor, dyn_feat: torch.Tensor
                ) -> Dict[str, torch.Tensor]:
        pooled = self.trunk_pool(beliefs, valid, text)              # [B, gru_hidden]
        side = self.dyn_proj(dyn_feat)                              # [B, dyn_hidden]
        joint = self.fusion(torch.cat([pooled, side], dim=-1))      # [B, 128]
        return {h: self.heads[h](joint).squeeze(-1) for h in self.head_names}

    def warm_start_from_pomdp(self, pomdp_state: Dict[str, torch.Tensor]):
        """Copy in_proj / text_proj / gru / attn weights from POMDPTemporalHead."""
        own = self.state_dict()
        copied = []
        for k, v in pomdp_state.items():
            if k in own and own[k].shape == v.shape:
                own[k] = v.clone()
                copied.append(k)
        self.load_state_dict(own)
        logger.info(f"warm-started {len(copied)} trunk params from POMDP")


# ─── train + eval loops ─────────────────────────────────────────────────────

def evaluate(model: LKAlertBDHead, ds: CacheDataset, head_names: List[str],
              device: torch.device, batch_size: int = 128) -> Dict:
    from sklearn.metrics import average_precision_score, roc_auc_score
    model.eval()
    loader = DataLoader(ds, batch_size=batch_size, shuffle=False,
                        collate_fn=collate)
    preds: Dict[str, List[float]] = {h: [] for h in head_names}
    labels: Dict[str, List[float]] = {h: [] for h in head_names}
    with torch.no_grad():
        for b in loader:
            dyn = build_features(b["belief"].to(device), b["valid"].to(device),
                                  b["tta_mean"].to(device),
                                  b["tta_var"].to(device))
            out = model(b["belief"].to(device), b["valid"].to(device),
                        b["text"].to(device), dyn["pooled"])
            for h in head_names:
                preds[h].extend(torch.sigmoid(out[h]).cpu().tolist())
                if f"y_{h}" in b:
                    labels[h].extend(b[f"y_{h}"].tolist())
    metrics: Dict[str, Dict] = {}
    for h in head_names:
        if not labels[h]:
            continue
        y = np.asarray(labels[h]); p = np.asarray(preds[h])
        if y.min() == y.max():
            metrics[h] = {"ap": 0.0, "auc": 0.0, "n": int(y.size),
                          "n_pos": int(y.sum())}
            continue
        metrics[h] = {
            "ap":     float(average_precision_score(y, p)),
            "auc":    float(roc_auc_score(y, p)),
            "n":      int(y.size),
            "n_pos":  int(y.sum()),
        }
    return metrics


def train_one_seed(args) -> Dict:
    torch.manual_seed(args.seed)
    np.random.seed(args.seed)
    random.seed(args.seed)

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

    # 1) labels
    train_lbl = derive_labels(DIAG_DIR / _diag_filename(args.train_cache),
                                ids := torch.load(
                                    CACHE_DIR / f"{args.train_cache}.pt",
                                    weights_only=False,
                                    map_location="cpu")["meta"]["ids"])
    head_names = usable_heads(train_lbl, args.head_names)
    if not head_names:
        raise SystemExit("no usable heads")
    logger.info(f"usable heads: {head_names}")
    for h in head_names:
        y = train_lbl[h]
        logger.info(f"  {h}: n_pos={int(y.sum())}  n_neg={int(y.size - y.sum())}")

    # 2) datasets
    train_ds = CacheDataset(CACHE_DIR / f"{args.train_cache}.pt",
                              DIAG_DIR / _diag_filename(args.train_cache),
                              head_names)
    val_caches: Dict[str, CacheDataset] = {}
    for vc in args.val_caches:
        diag = DIAG_DIR / _diag_filename(vc)
        diag_p = diag if diag.exists() else None
        val_caches[vc] = CacheDataset(CACHE_DIR / f"{vc}.pt", diag_p, head_names)

    # 3) sampler weights to balance p_any
    y_any = train_lbl["p_any"]
    pos = (y_any == 1).sum(); neg = (y_any == 0).sum()
    weights = np.where(y_any == 1, 1.0 / max(pos, 1), 1.0 / max(neg, 1))
    sampler = torch.utils.data.WeightedRandomSampler(
        weights=weights, num_samples=len(weights), replacement=True)

    train_loader = DataLoader(train_ds, batch_size=args.batch_size,
                                sampler=sampler, collate_fn=collate,
                                num_workers=args.num_workers)

    # 4) model
    in_dim = train_ds.bf.shape[-1]
    fdim = feature_dim(in_dim, with_tta=True)
    model = LKAlertBDHead(in_dim=in_dim, proj_dim=args.proj_dim,
                            gru_hidden=args.gru_hidden,
                            dyn_dim=fdim, dyn_hidden=args.dyn_hidden,
                            dropout=args.dropout, head_names=head_names)

    # 5) warm-start trunk from POMDP best
    if args.warm_start:
        ck = torch.load(args.warm_start, weights_only=False, map_location="cpu")
        model.warm_start_from_pomdp(ck["head_state"])

    model.to(device)
    opt = torch.optim.AdamW(model.parameters(), lr=args.lr,
                              weight_decay=args.wd)

    # 6) train loop
    best = {"epoch": -1, "macro_ap": -1.0, "per_cache": {}}
    out_dir = Path(args.out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    for epoch in range(args.epochs):
        model.train()
        ep_loss = 0.0
        n = 0
        for b in train_loader:
            opt.zero_grad()
            dyn = build_features(b["belief"].to(device), b["valid"].to(device),
                                  b["tta_mean"].to(device),
                                  b["tta_var"].to(device))
            logits = model(b["belief"].to(device), b["valid"].to(device),
                            b["text"].to(device), dyn["pooled"])
            losses = []
            for h in head_names:
                y = b[f"y_{h}"].to(device)
                losses.append(F.binary_cross_entropy_with_logits(logits[h], y))
            # ordinal monotonic regulariser p_500 ≤ p_1000 ≤ p_1500 ≤ p_any
            order = ["p_500", "p_1000", "p_1500", "p_any"]
            present = [h for h in order if h in head_names]
            if args.ordinal_lambda > 0 and len(present) >= 2:
                with torch.no_grad():
                    pass
                ord_loss = 0.0
                for a, c in zip(present[:-1], present[1:]):
                    ord_loss = ord_loss + F.relu(
                        torch.sigmoid(logits[a]) - torch.sigmoid(logits[c])
                    ).mean()
                losses.append(args.ordinal_lambda * ord_loss)
            loss = sum(losses) / max(len(losses), 1)
            loss.backward()
            opt.step()
            ep_loss += float(loss.detach()) * b["belief"].shape[0]
            n += b["belief"].shape[0]

        # eval
        per_cache: Dict[str, Dict] = {}
        macro_ap = 0.0
        for vc, ds in val_caches.items():
            m = evaluate(model, ds, head_names, device, args.batch_size)
            per_cache[vc] = m
            if "p_any" in m:
                macro_ap += m["p_any"]["ap"]
        macro_ap /= max(1, len(val_caches))
        logger.info(f"ep {epoch:02d}  loss={ep_loss/max(n,1):.4f}  "
                    f"macro p_any AP={macro_ap:.4f}")
        for vc, m in per_cache.items():
            for h, mh in m.items():
                logger.info(f"  {vc}/{h}: AP={mh['ap']:.4f} AUC={mh['auc']:.4f} "
                            f"n_pos={mh['n_pos']}/{mh['n']}")
        if macro_ap > best["macro_ap"]:
            best = {"epoch": epoch, "macro_ap": float(macro_ap),
                    "per_cache": per_cache,
                    "head_state": model.state_dict(),
                    "args": vars(args), "head_names": head_names}
            torch.save(best, out_dir / "best.pt")
            logger.info(f"  -> saved best.pt @ macro_ap={macro_ap:.4f}")
    return best


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--train_cache", default="nexar_train_diag")
    ap.add_argument("--val_caches",  nargs="+",
                    default=["nexar_val", "dota_val", "dad_test", "dada_test"])
    ap.add_argument("--out_dir",     default="checkpoints/Policy/lkalert_bd_seed0")
    ap.add_argument("--epochs", type=int, default=30)
    ap.add_argument("--batch_size", type=int, default=64)
    ap.add_argument("--num_workers", type=int, default=2)
    ap.add_argument("--lr", type=float, default=3e-4)
    ap.add_argument("--wd", type=float, default=1e-4)
    ap.add_argument("--proj_dim", type=int, default=512)
    ap.add_argument("--gru_hidden", type=int, default=256)
    ap.add_argument("--dyn_hidden", type=int, default=64)
    ap.add_argument("--dropout", type=float, default=0.2)
    ap.add_argument("--ordinal_lambda", type=float, default=0.1)
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--head_names", nargs="+",
                    default=["p_any", "p_1500", "p_1000", "p_500",
                             "p_resolution_proxy", "p_ego"],
                    help="degenerate buckets are silently dropped")
    ap.add_argument("--warm_start", default=
                    "checkpoints/Policy/pomdp_head_qwen3vl4b_best_seed/best.pt",
                    help="POMDP best.pt to warm-start trunk")
    args = ap.parse_args()
    train_one_seed(args)


if __name__ == "__main__":
    main()