"""Score ALL VLAlert / LKAlert variants on benchmark/v1/val using a shared belief cache.

Runs each (danger_ckpt, policy_ckpt) combo through the sft_x_v3 cache
and writes per-tick PT files in the v1/val schema for downstream
aggregation.

Usage:
    python tools/score_v1_val_vlalert_all.py \
        --cache data/belief_cache_v2/sft_x_v3__v1_val.pt \
        --manifest eval_results/benchmark_v1_val/val_manifest.json

Output schema (each .pt in eval_results/benchmark_v1_val/per_tick/):
  Same as tools/score_v1_val_baselines.py — see that file for full schema.
"""
from __future__ import annotations
import argparse
import json
import logging
import sys
import time
from pathlib import Path

import torch
import torch.nn.functional as F
from tqdm import tqdm

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))
from lkalert.models.danger_head import DangerHead
from lkalert.models.policy_head_v2 import PolicyHeadV2

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


# ── Variant registry: (display_name, danger_ckpt, policy_ckpt, belief_slice_dim) ──
# All share sft_x_v3 backbone via the cache. `belief_slice_dim` is None (= use full
# 10240-d cache) or 2560 (= use only L32 = last 2560 dims, for c1_lastonly variants).
VARIANTS = [
    # Headline / paper-facing
    ("VLAlert-X",              "checkpoints/danger_v3_hazard/best.pt", "checkpoints/policy_v3_strong/best.pt", None),
    ("VLAlert-X-v2",           "checkpoints/danger_v3_hazard/best.pt", "checkpoints/policy_v3_strong_v2/best.pt", None),
    # RL variants (head-DPO/KTO/PPO; VLM frozen)
    ("VLAlert-X+Head-DPO",     "checkpoints/danger_v3_hazard/best.pt", "checkpoints/policy_v3_head_dpo/best.pt", None),
    ("VLAlert-X+Head-KTO",     "checkpoints/danger_v3_hazard/best.pt", "checkpoints/policy_v3_head_kto/best.pt", None),
    ("VLAlert-X+Head-PPO",     "checkpoints/danger_v3_hazard/best.pt", "checkpoints/policy_v3_head_ppo/best.pt", None),
    # Closed-loop / adaptive variants
    ("VLAlert-X+Adaptive",     "checkpoints/danger_v3_hazard/best.pt", "checkpoints/policy_v3_adaptive/best.pt", None),
    ("VLAlert-X+Adaptive-DPO", "checkpoints/danger_v3_hazard/best.pt", "checkpoints/policy_v3_adaptive_dpo/best.pt", None),
    ("VLAlert-X+Adaptive-DPO-v2", "checkpoints/danger_v3_hazard/best.pt", "checkpoints/policy_v3_adaptive_dpo_v2/best.pt", None),
    ("VLAlert-X+Adaptive-relabel", "checkpoints/danger_v3_hazard/best.pt", "checkpoints/policy_v3_adaptive_relabel/best.pt", None),
    # Layer-ablation 5-seed (c1_lastonly: only L32, 2560-d belief). Paired with
    # the canonical v3_strong PolicyHead since these ckpts are DangerHead-only.
    ("VLAlert-X+c1-seed1",     "checkpoints/layer_ablation_v2/c1_lastonly_seed1/best.pt", "checkpoints/policy_v3_strong/best.pt", 2560),
    ("VLAlert-X+c1-seed2",     "checkpoints/layer_ablation_v2/c1_lastonly_seed2/best.pt", "checkpoints/policy_v3_strong/best.pt", 2560),
    ("VLAlert-X+c1-seed3",     "checkpoints/layer_ablation_v2/c1_lastonly_seed3/best.pt", "checkpoints/policy_v3_strong/best.pt", 2560),
    ("VLAlert-X+c1-seed4",     "checkpoints/layer_ablation_v2/c1_lastonly_seed4/best.pt", "checkpoints/policy_v3_strong/best.pt", 2560),
    ("VLAlert-X+c1-seed5",     "checkpoints/layer_ablation_v2/c1_lastonly_seed5/best.pt", "checkpoints/policy_v3_strong/best.pt", 2560),
    # v4 experimental adaptive (2 seeds)
    ("VLAlert-X+v4-Adaptive-seed0", "checkpoints/danger_v3_hazard/best.pt", "checkpoints/policy_v4_adaptive/seed0/best.pt", None),
    ("VLAlert-X+v4-Adaptive-seed1", "checkpoints/danger_v3_hazard/best.pt", "checkpoints/policy_v4_adaptive/seed1/best.pt", None),
    # ── Legacy v2-M10 (5 seeds) — paired with danger_v2/seed2 (their training pairing)
    # Cross-backbone: scored on sft_x_v3 cache; treats v3 belief features as input.
    # User accepts architecture drift ("思想相同") so we report honestly.
    ("VLAlert-v2-M10-seed0",   "checkpoints/danger_v2/seed2/best.pt", "checkpoints/policy_v2/seed0/best.pt", None),
    ("VLAlert-v2-M10-seed1",   "checkpoints/danger_v2/seed2/best.pt", "checkpoints/policy_v2/seed1/best.pt", None),
    ("VLAlert-v2-M10-seed2",   "checkpoints/danger_v2/seed2/best.pt", "checkpoints/policy_v2/seed2/best.pt", None),
    ("VLAlert-v2-M10-seed3",   "checkpoints/danger_v2/seed2/best.pt", "checkpoints/policy_v2/seed3/best.pt", None),
    ("VLAlert-v2-M10-seed4",   "checkpoints/danger_v2/seed2/best.pt", "checkpoints/policy_v2/seed4/best.pt", None),
    # ── Legacy v3 family (3 CE + 3 focord seeds) — paired with danger_v2/seed2
    ("VLAlert-v3-CE-seed0",       "checkpoints/danger_v2/seed2/best.pt", "checkpoints/policy_v3_full/ce_seed0/best.pt", None),
    ("VLAlert-v3-CE-seed1",       "checkpoints/danger_v2/seed2/best.pt", "checkpoints/policy_v3_full/ce_seed1/best.pt", None),
    ("VLAlert-v3-CE-seed2",       "checkpoints/danger_v2/seed2/best.pt", "checkpoints/policy_v3_full/ce_seed2/best.pt", None),
    ("VLAlert-v3-Focord-seed0",   "checkpoints/danger_v2/seed2/best.pt", "checkpoints/policy_v3_full/focord_seed0/best.pt", None),
    ("VLAlert-v3-Focord-seed1",   "checkpoints/danger_v2/seed2/best.pt", "checkpoints/policy_v3_full/focord_seed1/best.pt", None),
    ("VLAlert-v3-Focord-seed2",   "checkpoints/danger_v2/seed2/best.pt", "checkpoints/policy_v3_full/focord_seed2/best.pt", None),
    # (SFT-argmax is handled separately via tools/eval_sft_argmax_baseline.py + converter)
]


@torch.no_grad()
def score_one(name: str, danger_ckpt: Path, policy_ckpt: Path,
              cache: dict, val_manifest_samples: list,
              batch_size: int, device: torch.device,
              prev_action: int = 3,
              belief_slice_dim: int = None) -> dict:
    """Run one (danger, policy) combo on the shared cache.

    Returns extended schema including: raw_logits, scores_3class, scores_binary,
    danger_per_frame, danger_clip, perception_summary, first_fire_tta, lead_time.
    """
    print(f"\n══════════ {name} ══════════")
    print(f"  danger: {danger_ckpt}")
    print(f"  policy: {policy_ckpt}")
    if not danger_ckpt.exists():
        print(f"  [skip] danger ckpt missing")
        return None
    if not policy_ckpt.exists():
        print(f"  [skip] policy ckpt missing")
        return None

    belief_full = cache["belief_content"].float()  # [N, 8, D_belief_full]
    # Slice belief to last K dims if requested (for c1_lastonly variants which
    # were trained on L32-only). Cache stacks layers [L20, L24, L28, L32] with
    # L32 = last 2560 dims.
    if belief_slice_dim is not None:
        belief = belief_full[:, :, -belief_slice_dim:].contiguous()
        print(f"  belief sliced to last {belief_slice_dim} dims (c1_lastonly variant)")
    else:
        belief = belief_full
    policy = cache["policy_position"].float()   # [N, 8, D_policy]
    valid  = cache["valid_frames"]               # [N, 8]
    N = belief.shape[0]

    # ── load heads ──
    ck_d = torch.load(danger_ckpt, weights_only=False, map_location="cpu")
    if ck_d["in_dim"] != belief.shape[-1]:
        print(f"  [skip] danger ckpt in_dim={ck_d['in_dim']} != belief dim={belief.shape[-1]}")
        return None
    danger = DangerHead(in_dim=ck_d["in_dim"]).to(device)
    # strict=False tolerates extra modules (e.g., hazard_head sub-classifier
    # in danger_v3_hazard ckpt) that aren't part of the base DangerHead API.
    missing, unexpected = danger.load_state_dict(ck_d["model"], strict=False)
    if unexpected:
        print(f"  [info] unexpected keys (ignored): {len(unexpected)}")
    if missing:
        print(f"  [warn] missing keys: {missing[:3]}")
        return None
    danger.eval()

    ck_p = torch.load(policy_ckpt, weights_only=False, map_location="cpu")
    try:
        policy_head = PolicyHeadV2(
            policy_dim=ck_p["policy_dim"],
            perception_dim_per_query=ck_p["perception_dim_per_query"],
            k_queries=ck_p["k_queries"],
        ).to(device)
        # Legacy v2/v3 ckpts used a single Sequential `fuse` (Linear,GELU,Dropout,Linear);
        # the new PolicyHeadV2 splits it into `fuse_pre` (first 3 modules) and
        # `cls_head` (final Linear). Remap keys for backward compat.
        sd = ck_p["model"]
        if any(k.startswith("fuse.") for k in sd):
            remapped = {}
            for k, v in sd.items():
                if k.startswith("fuse.0."):       # Linear → fuse_pre.0
                    remapped["fuse_pre.0." + k[len("fuse.0."):]] = v
                elif k.startswith("fuse.3."):     # final Linear → cls_head
                    remapped["cls_head." + k[len("fuse.3."):]] = v
                else:
                    remapped[k] = v
            sd = remapped
            print("  [info] remapped legacy fuse.{0,3} → fuse_pre.0 + cls_head")
        policy_head.load_state_dict(sd, strict=False)
        policy_head.eval()
    except (KeyError, RuntimeError) as e:
        print(f"  [skip] policy ckpt incompatible: {e}")
        return None

    # ── infer ──
    raw_logits_out = torch.zeros(N, 3, dtype=torch.float32)
    danger_pf_out  = torch.zeros(N, 8, dtype=torch.float32)
    danger_clip_out = torch.zeros(N, dtype=torch.float32)
    # perception_summary shape is [B, K, hidden]; allocate when we know K, hidden
    perception_out = None
    prev_act = torch.full((batch_size,), prev_action, dtype=torch.long, device=device)
    t0 = time.time()
    for i in tqdm(range(0, N, batch_size), ncols=80, desc=f"infer {name}"):
        end = min(N, i + batch_size)
        b_belief = belief[i:end].to(device, non_blocking=True)
        b_policy = policy[i:end].to(device, non_blocking=True)
        b_valid  = valid[i:end].to(device, non_blocking=True)
        cur_bs = end - i
        d_out = danger(b_belief, valid_frames=b_valid)
        perc = d_out["perception_summary"]
        danger_pf = d_out["per_frame"]
        danger_clip = d_out["clip"]
        if perception_out is None:
            K, H = perc.shape[1], perc.shape[2]
            perception_out = torch.zeros(N, K, H, dtype=torch.float32)
        perception_out[i:end] = perc.float().cpu()
        danger_pf_out[i:end] = danger_pf.float().cpu()
        danger_clip_out[i:end] = danger_clip.float().cpu()
        prev = prev_act[:cur_bs]
        logits = policy_head(b_policy, perc, danger_pf, prev, valid_frames=b_valid)
        raw_logits_out[i:end] = logits.float().cpu()
    print(f"  inference: {time.time()-t0:.0f}s")

    s3c = F.softmax(raw_logits_out, dim=-1)
    s_bin = s3c[:, 2].clone()

    # ── pull metadata from cache (correct per-tick) ──
    # Cache stores 'ids' = synthetic ("v1val_006901") and 'video_id' = real
    # ("nexar_00002"). Use video_id for matching with the val_manifest.
    ids       = list(cache.get("video_id", cache["ids"]))
    sources   = list(cache.get("source", [""] * N))
    raw_cats  = list(cache.get("category", [""] * N))
    ttas      = cache.get("tick_tta_raw", torch.full((N,), -1.0)).float()

    # ── pull GT label directly from cache (cache stores tick_action) ──
    # The cache ALREADY has correct per-tick labels (tick_action) and metadata.
    # We supplement only fps/n_frames/tick_idx from manifest if available.
    tick_action_cache = cache.get("tick_action", torch.zeros(N, dtype=torch.long))
    label_out = tick_action_cache.tolist()

    # Lookup table for fps/n_frames/tick_idx (manifest only used for these).
    # Match by video_id; for each video, ticks are in cache order so we use
    # appearance-order to assign tick_idx within a video.
    manifest_by_vid = {}
    for s in val_manifest_samples:
        manifest_by_vid.setdefault(s["video_id"], []).append(s)

    fi_out, fps_out, nframes_out, tidx_out, cat_hf_out = [], [], [], [], []
    n_empty = 0
    vid_seen_count: dict = {}
    for i in range(N):
        vid = ids[i]
        if not vid:
            # Cache extractor failed for this tick (e.g., DoTA frame-folder bug)
            n_empty += 1
            fi_out.append([0] * 8)
            fps_out.append(30.0)
            nframes_out.append(0)
            tidx_out.append(0)
            cat_hf_out.append("?")
            continue
        ms = manifest_by_vid.get(vid, [])
        k = vid_seen_count.get(vid, 0)
        m = ms[k] if k < len(ms) else (ms[0] if ms else None)
        vid_seen_count[vid] = k + 1
        if m is None:
            fi_out.append([0] * 8); fps_out.append(30.0); nframes_out.append(0); tidx_out.append(0); cat_hf_out.append("?")
            continue
        fi_out.append(list(m["frame_indices"]))
        fps_out.append(float(m["fps"]))
        nframes_out.append(int(m["n_frames"]))
        tidx_out.append(int(m.get("tick_idx", k)))
        cat_hf_out.append(m["category"])
    if n_empty:
        print(f"  [warn] {n_empty} ticks have empty cache entries (likely DoTA frame-folder failures)")

    # ── post-hoc derive first_fire_tta + lead_time per tick ──
    # For each video, find the FIRST tick where s_bin >= 0.5; record its
    # tta_raw and compute lead_time = max(0, that_tick.tta_raw).
    # Each tick stores (first_fire_tta, lead_time) inherited from the
    # video's first-fire tick (or NaN if never fires).
    tick_label_t = torch.tensor(label_out, dtype=torch.long)
    fps_t = torch.tensor(fps_out, dtype=torch.float)
    tidx_t = torch.tensor(tidx_out, dtype=torch.long)
    first_fire_tta_out = torch.full((N,), float("nan"), dtype=torch.float)
    lead_time_out = torch.full((N,), float("nan"), dtype=torch.float)
    # Group ticks by video_id
    from collections import defaultdict as _dd
    by_video = _dd(list)
    for i in range(N):
        by_video[ids[i]].append(i)
    for vid, idxs in by_video.items():
        # sort by tick_idx ascending
        idxs_sorted = sorted(idxs, key=lambda j: tidx_t[j].item())
        fired = False
        for j in idxs_sorted:
            if not fired and s_bin[j].item() >= 0.5:
                first_fire_tta_out[j] = ttas[j].item()
                # lead_time = tta_raw at first fire (positive = before event)
                lead_time_out[j] = max(0.0, float(ttas[j].item()))
                fired = True
        # Distribute first-fire info to all ticks of this video for convenience
        if fired:
            for j in idxs_sorted:
                if torch.isnan(first_fire_tta_out[j]):
                    first_fire_tta_out[j] = first_fire_tta_out[idxs_sorted[0]] if not torch.isnan(first_fire_tta_out[idxs_sorted[0]]) else float("nan")

    # Compose schema-conforming output (EXTENDED)
    out = {
        # metadata
        "method":              name,
        "ckpt":                str(policy_ckpt),
        "danger_ckpt":         str(danger_ckpt),
        "belief_slice_dim":    belief_slice_dim,
        "manifest":            "eval_results/benchmark_v1_val/val_manifest.json",
        "n_ticks":             int(N),
        # tick-level identifiers
        "ids":                 ids,
        "source":              sources,
        "category":            cat_hf_out,
        "raw_category":        raw_cats,
        "frame_indices":       torch.tensor(fi_out, dtype=torch.long),
        "tta_raw":             ttas,
        "fps":                 fps_t,
        "n_frames":            torch.tensor(nframes_out, dtype=torch.long),
        "tick_idx":            tidx_t,
        "tick_label":          tick_label_t,
        # primary scores
        "raw_logits":          raw_logits_out,
        "scores_3class":       s3c,
        "scores_binary":       s_bin,
        # NEW intermediate variables (for calibration / re-analysis)
        "danger_per_frame":    danger_pf_out,
        "danger_clip":         danger_clip_out,
        "perception_summary":  perception_out,
        "prev_action_used":    torch.full((N,), prev_action, dtype=torch.long),
        "first_fire_tta":      first_fire_tta_out,
        "lead_time":           lead_time_out,
    }
    return out


def report_brief(out: dict):
    import numpy as np
    from sklearn.metrics import average_precision_score, roc_auc_score
    y_true = out["tick_label"].numpy()
    y_alert = (y_true == 2).astype(np.int64)
    scores = out["scores_binary"].numpy()
    try:
        ap = average_precision_score(y_alert, scores)
        auc = roc_auc_score(y_alert, scores) if 0 < y_alert.sum() < len(y_alert) else float("nan")
    except Exception:
        ap = auc = float("nan")
    print(f"  binary AP={ap:.4f}  AUROC={auc:.4f}  n_pos={y_alert.sum()}")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--cache",    type=Path,
                    default=ROOT / "data/belief_cache_v2/sft_x_v3__v1_val.pt")
    ap.add_argument("--manifest", type=Path,
                    default=ROOT / "eval_results/benchmark_v1_val/val_manifest.json")
    ap.add_argument("--out_dir",  type=Path,
                    default=ROOT / "eval_results/benchmark_v1_val/per_tick")
    ap.add_argument("--batch_size", type=int, default=64)
    ap.add_argument("--prev_action", type=int, default=3)
    ap.add_argument("--variants", nargs="+", default=None,
                    help="Subset of variant names to score (default: all)")
    args = ap.parse_args()

    args.out_dir.mkdir(parents=True, exist_ok=True)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"[device] {device}")
    print(f"[cache]  {args.cache}")
    print(f"[manifest] {args.manifest}")

    if not args.cache.exists():
        print(f"[err] cache not found — wait for extraction to finish first")
        return

    # Load cache once (reused across variants)
    print(f"[load] cache ...")
    cache = torch.load(args.cache, weights_only=False, map_location="cpu")
    val_doc = json.loads(args.manifest.read_text())
    val_samples = val_doc["samples"]

    # Optional filter
    to_run = VARIANTS
    if args.variants:
        to_run = [v for v in VARIANTS if v[0] in args.variants]

    for variant in to_run:
        # Support both (name, danger, policy) and (name, danger, policy, slice_dim)
        if len(variant) == 4:
            name, dpath, ppath, slice_dim = variant
        else:
            name, dpath, ppath = variant
            slice_dim = None
        try:
            out = score_one(name, ROOT / dpath, ROOT / ppath,
                            cache, val_samples, args.batch_size, device,
                            args.prev_action, belief_slice_dim=slice_dim)
            if out is None:
                continue
            slug = name.lower().replace("+", "_").replace(" ", "_").replace("-", "_")
            out_path = args.out_dir / f"{slug}.pt"
            torch.save(out, out_path)
            print(f"  [save] {out_path}")
            report_brief(out)
        except Exception as e:
            print(f"  [error scoring {name}]: {e}")
            import traceback; traceback.print_exc()


if __name__ == "__main__":
    main()