#!/usr/bin/env python
"""Demo comparison pipeline: score all videos with multiple models, generate viz videos.

Models (scored in backbone order to maximise GPU reuse):
  1. BADAS (V-JEPA2) — 16-frame sliding window
  2. VLAlert-v3 — sft_x_v3 + danger_v3 + policy_v3_strong
  3. VLAlert-v2 — sft_x_v2 + danger_v2 + policy_v2_full (5-seed ensemble)
  4. VLAlert-X  — sft_x_v2 + VLAlertXHead (5-seed ensemble, narrow window)
  5. VLAlert-M10 — qwen3vl4b_cot_belief_perframe + M10 head (5-seed ensemble)

Pipeline:
  Phase 1: Extract frames (already done → demo/compare_frames/)
  Phase 2: Score all videos model-by-model (one VLM backbone at a time)
  Phase 3: Generate comparison videos (left=frame, right=score+action)

Usage:
    python tools/demo_compare_pipeline.py [--models v3,X,v2,M10] [--only VIDEO]
"""
from __future__ import annotations
import argparse, cv2, gc, json, logging, sys, time
from pathlib import Path
import numpy as np
import torch
from PIL import Image
from tqdm import tqdm

ROOT = Path("PROJECT_ROOT")
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

# ─── Conv3d → Linear patch for Qwen3-VL (64× speedup on Blackwell) ───
import torch.nn as nn
from transformers.models.qwen3_vl.modeling_qwen3_vl import Qwen3VLVisionPatchEmbed

def _fast_patch_embed_forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
    target_dtype = self.proj.weight.dtype
    if isinstance(self.proj, nn.Conv3d):
        conv = self.proj
        out_dim = conv.out_channels
        in_dim = (conv.in_channels * conv.kernel_size[0]
                  * conv.kernel_size[1] * conv.kernel_size[2])
        w_flat = conv.weight.detach().reshape(out_dim, in_dim).contiguous()
        bias = conv.bias.detach().clone() if conv.bias is not None else None
        new_proj = nn.Linear(in_dim, out_dim, bias=bias is not None)
        new_proj.weight.data.copy_(w_flat)
        if bias is not None:
            new_proj.bias.data.copy_(bias)
        new_proj.to(device=conv.weight.device, dtype=conv.weight.dtype)
        self.proj = new_proj
    if hidden_states.dim() > 2 or hidden_states.shape[-1] != self.proj.in_features:
        hidden_states = hidden_states.reshape(-1, self.proj.in_features)
    return self.proj(hidden_states.to(dtype=target_dtype))

Qwen3VLVisionPatchEmbed.forward = _fast_patch_embed_forward
FRAMES_DIR = ROOT / "demo/compare_frames"
OUT_DIR = ROOT / "demo/compare_results"
OUT_DIR.mkdir(exist_ok=True)

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

# ─── BADAS config ───
BADAS_REPO = Path("~/.cache/huggingface/hub/models--nexar-ai--badas-open/"
                   "snapshots/8fda93711e79d72401b0a4efc151b56455885cd2")
BADAS_MODEL = "facebook/vjepa2-vitl-fpc16-256-ssv2"
BADAS_CKPT = str(BADAS_REPO / "weights" / "badas_open.pth")

# ─── VLAlert configs ───
SFT_V3 = ROOT / "checkpoints/sft_x_v3/best"
SFT_V2 = ROOT / "checkpoints/sft_x_v2/best"
SFT_B0 = ROOT / "checkpoints/VLA/qwen3vl4b_cot_belief_perframe/best"
DANGER_V3 = ROOT / "checkpoints/danger_v3_hazard/best.pt"
DANGER_V2 = ROOT / "checkpoints/danger_v2/seed2/best.pt"
POLICY_V3 = ROOT / "checkpoints/policy_v3_strong/best.pt"
POLICY_V2_SEEDS = [ROOT / f"checkpoints/policy_v2_full/seed{s}/best.pt" for s in range(5)]
POLICY_X_SEEDS = [ROOT / f"checkpoints/policy_x_L4_bal_seed{s}/best.pt" for s in range(5)]
M10_SEEDS = [ROOT / f"checkpoints/Policy/m10_qwen3vl4b_seed{s}/best/policy_head.pt" for s in range(5)]
BASE_MODEL = ROOT / "models/Qwen3-VL-4B-Instruct"

# ─── Qwen2.5-VL-3B config ───
BASE_MODEL_Q25 = ROOT / "models/Qwen2.5-VL-3B-Instruct"
SFT_Q25_LORA = ROOT / "checkpoints/sft/sft_qwen25vl3b_lora_resume/best/vlm_lora"
TTA_HEAD_Q25 = ROOT / "checkpoints/sft/sft_qwen25vl3b_lora_resume/best/tta_head.pt"


def free_gpu():
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()


import os
VLM_MAX_DIM = int(os.environ.get("VLM_MAX_DIM", "0"))

def load_frames(video_dir: Path, indices: list[int]) -> list[Image.Image]:
    """Load PIL frames by index from extracted jpg folder."""
    out = []
    for fi in indices:
        for fmt in [f"{fi:06d}.jpg", f"{fi:05d}.jpg", f"{fi:04d}.jpg",
                    f"{fi:03d}.jpg", f"{fi}.jpg"]:
            p = video_dir / fmt
            if p.exists():
                img = Image.open(p).convert("RGB")
                if VLM_MAX_DIM > 0 and max(img.size) > VLM_MAX_DIM:
                    r = VLM_MAX_DIM / max(img.size)
                    nw = max(int(img.width * r) // 28 * 28, 28)
                    nh = max(int(img.height * r) // 28 * 28, 28)
                    img = img.resize((nw, nh), Image.BILINEAR)
                out.append(img)
                break
        else:
            if out:
                out.append(out[-1])
            else:
                out.append(Image.new("RGB", (640, 360)))
    return out


def uniform_indices(start, end, n):
    if end <= start: return [start] * n
    return np.linspace(start, end, n).round().astype(int).tolist()


# ═══════════════════════════════════════════════════════════════
#  BADAS scorer
# ═══════════════════════════════════════════════════════════════
class BADASScorer:
    def __init__(self):
        sys.path.insert(0, str(BADAS_REPO / "src"))
        import train.video_training  # noqa
        from models.vjepa import VJEPAModel
        logger.info("[BADAS] loading V-JEPA2...")
        self.vjepa = VJEPAModel(
            model_name=BADAS_MODEL, checkpoint_path=BADAS_CKPT,
            frame_count=16, img_size=224, window_stride=1,
            target_fps=8.0, use_sliding_window=False)
        self.vjepa.load()
        self.device = self.vjepa.device

    @torch.no_grad()
    def score_tick(self, frames_16: list[Image.Image]) -> float:
        proc = self.vjepa.processor(videos=[frames_16], return_tensors="pt")
        key = "pixel_values_videos" if "pixel_values_videos" in proc else "pixel_values"
        video = proc[key].to(self.device)
        if video.dim() == 4: video = video.unsqueeze(0)
        with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
            out = self.vjepa.model(video)
        logits = out.float() / 2.0
        return float(torch.softmax(logits, dim=1)[0, 1].cpu())

    def score_video(self, video_dir: Path, n_frames: int, fps: float, **kw) -> list[dict]:
        """Score at 1Hz ticks."""
        results = []
        tick_interval = max(1, int(fps))
        for tick_frame in range(0, n_frames, tick_interval):
            end = min(tick_frame, n_frames - 1)
            start = max(0, end - 15)
            indices = uniform_indices(start, end, 16)
            frames = load_frames(video_dir, indices)
            p = self.score_tick(frames)
            action = "ALERT" if p > 0.5 else ("OBSERVE" if p > 0.07 else "SILENT")
            results.append({"frame": tick_frame, "t": tick_frame / fps,
                             "p_alert": p, "action": action})
        return results


# ═══════════════════════════════════════════════════════════════
#  VLAlert scorer (v3 or X)
# ═══════════════════════════════════════════════════════════════
class VLAlertScorer:
    def __init__(self, sft_path, danger_path, policy_paths, name="VLAlert"):
        self.name = name
        self.device = "cuda" if torch.cuda.is_available() else "cpu"

        # Load DangerHead
        from lkalert.models.danger_head import DangerHead
        ck = torch.load(danger_path, weights_only=False, map_location="cpu")
        self.danger = DangerHead(in_dim=ck["in_dim"],
                                  n_hazards=int(ck.get("n_hazards", 0) or 0)).to(self.device)
        self.danger.load_state_dict(ck["model"])
        self.danger.eval()

        # Load PolicyHead(s)
        from lkalert.models.policy_head_v2 import PolicyHeadV2
        self.policies = []
        for pp in policy_paths:
            pk = torch.load(pp, weights_only=False, map_location="cpu")
            policy = PolicyHeadV2(
                policy_dim=pk.get("policy_dim", pk.get("in_dim", 2560)),
                perception_dim_per_query=pk.get("perception_dim_per_query", 512),
                k_queries=pk.get("k_queries", 4),
            ).to(self.device)
            sd = pk["model"]
            mapped = {}
            for k, v in sd.items():
                nk = k.replace("fuse.0.", "fuse_pre.0.").replace("fuse.3.", "cls_head.")
                mapped[nk] = v
            policy.load_state_dict(mapped, strict=False)
            policy.eval()
            self.policies.append(policy)

        # VLM belief cache (lazily populated per video)
        self.belief_cache = None
        self.sft_path = sft_path
        self.vlm_loaded = False
        logger.info(f"[{name}] danger + {len(self.policies)} policy heads loaded")

    def _ensure_vlm(self):
        if self.vlm_loaded: return
        logger.info(f"[{self.name}] loading VLM from {self.sft_path}...")
        from transformers import AutoProcessor, AutoModelForImageTextToText
        from peft import PeftModel
        from training.VLA.cot_belief_dataset_v2 import ALL_SPECIAL, BELIEF_OPEN, BELIEF_CLOSE, build_chat_v2

        self.processor = AutoProcessor.from_pretrained(BASE_MODEL, trust_remote_code=True)
        self.processor.tokenizer.add_special_tokens({"additional_special_tokens": ALL_SPECIAL})
        self.processor.tokenizer.padding_side = "right"

        base = AutoModelForImageTextToText.from_pretrained(
            BASE_MODEL, torch_dtype=torch.bfloat16, trust_remote_code=True)
        base.resize_token_embeddings(len(self.processor.tokenizer))
        self.vlm = PeftModel.from_pretrained(base, self.sft_path).to(self.device)
        self.vlm.eval()

        self.belief_open_id = self.processor.tokenizer.convert_tokens_to_ids(BELIEF_OPEN)
        self.belief_close_id = self.processor.tokenizer.convert_tokens_to_ids(BELIEF_CLOSE)
        self.belief_layers = [20, 24, 28, 32]
        self.policy_layer = 33
        self.build_chat = build_chat_v2
        self.vlm_loaded = True
        logger.info(f"[{self.name}] VLM loaded")

    @torch.no_grad()
    def extract_belief_batch(self, frames_batch: list[list[Image.Image]]):
        """Batch extract beliefs. frames_batch: list of N × [8 PIL images].
        Returns belief [N,8,10240], policy [N,8,2560], valid [N,8].
        """
        self._ensure_vlm()
        from training.VLA.cot_belief_dataset_v2 import SYSTEM_PROMPT_V2, USER_PROMPT_V2

        N = len(frames_batch)
        texts = []
        all_images = []
        for frames_8 in frames_batch:
            user_content = [{"type": "image", "image": img} for img in frames_8]
            user_content.append({"type": "text", "text": USER_PROMPT_V2})
            msgs = [
                {"role": "system", "content": [{"type": "text", "text": SYSTEM_PROMPT_V2}]},
                {"role": "user", "content": user_content},
            ]
            texts.append(self.processor.apply_chat_template(
                msgs, add_generation_prompt=True, tokenize=False))
            all_images.extend(frames_8)

        inputs = self.processor(text=texts, images=all_images, return_tensors="pt",
                                padding=True).to(self.device)

        with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
            out = self.vlm(**inputs, output_hidden_states=True, return_dict=True)
        hs_tuple = out.hidden_states
        D = hs_tuple[self.belief_layers[0]].shape[-1]

        belief = torch.zeros(N, 8, len(self.belief_layers) * D, dtype=torch.float16)
        policy = torch.zeros(N, 8, D, dtype=torch.float16)
        valid = torch.zeros(N, 8, dtype=torch.bool)

        for i in range(N):
            ids = inputs["input_ids"][i]
            open_pos = (ids == self.belief_open_id).nonzero(as_tuple=False).flatten().tolist()
            close_pos = (ids == self.belief_close_id).nonzero(as_tuple=False).flatten().tolist()
            n_blocks = min(len(open_pos), len(close_pos), 8)
            for f in range(n_blocks):
                o, c = open_pos[f], close_pos[f]
                if c <= o + 1:
                    continue
                parts = [hs_tuple[L][i, o+1:c].mean(dim=0).to(torch.float16)
                         for L in self.belief_layers]
                belief[i, f] = torch.cat(parts, dim=-1).cpu()
                policy[i, f] = hs_tuple[self.policy_layer][i, c].to(torch.float16).cpu()
                valid[i, f] = True

        del out, hs_tuple, inputs
        torch.cuda.empty_cache()
        return belief, policy, valid

    @torch.no_grad()
    def score_heads_batch(self, belief, policy_pos, valid):
        """Run DangerHead + PolicyHeads on batch. Returns list of (p_alert, p_obs, action, clip_danger)."""
        b = belief.to(self.device, dtype=torch.float32)
        v = valid.to(self.device)
        d_out = self.danger(b, valid_frames=v)
        perc = d_out["perception_summary"]
        dang = d_out["per_frame"]
        pp = policy_pos.to(self.device, dtype=torch.float32)
        N = b.shape[0]
        prev = torch.full((N,), 3, device=self.device, dtype=torch.long)

        probs_list = []
        for pol in self.policies:
            logits = pol(pp, perc, dang, prev, valid_frames=v)
            probs_list.append(torch.softmax(logits, dim=-1))
        avg = torch.stack(probs_list).mean(dim=0)

        results = []
        for i in range(N):
            p_alert = float(avg[i, 2].cpu())
            p_obs = float(avg[i, 1].cpu())
            act_idx = int(avg[i].argmax().cpu())
            action = ["SILENT", "OBSERVE", "ALERT"][act_idx]
            results.append((p_alert, p_obs, action, float(d_out["clip"][i].cpu())))
        return results

    def score_video(self, video_dir: Path, n_frames: int, fps: float,
                    batch_size: int = 2) -> list[dict]:
        tick_interval = max(1, int(fps))
        tick_frames = list(range(0, n_frames, tick_interval))

        all_frame_sets = []
        for tf in tick_frames:
            end = min(tf + 7, n_frames - 1)
            start = max(0, end - 7)
            indices = list(range(start, end + 1))
            while len(indices) < 8:
                indices = [indices[0]] + indices
            all_frame_sets.append(load_frames(video_dir, indices[:8]))

        results = []
        for bi in tqdm(range(0, len(tick_frames), batch_size),
                       desc=f"{self.name}", ncols=80, leave=False):
            batch_frames = all_frame_sets[bi:bi + batch_size]
            belief, policy_pos, valid = self.extract_belief_batch(batch_frames)
            head_results = self.score_heads_batch(belief, policy_pos, valid)
            for j, (p_alert, p_obs, action, clip_d) in enumerate(head_results):
                tf = tick_frames[bi + j]
                results.append({
                    "frame": tf, "t": tf / fps,
                    "p_alert": p_alert, "p_observe": p_obs,
                    "clip_danger": clip_d, "action": action,
                })
        return results

    def unload_vlm(self):
        if self.vlm_loaded:
            del self.vlm
            self.vlm_loaded = False
            free_gpu()
            logger.info(f"[{self.name}] VLM unloaded")


# ═══════════════════════════════════════════════════════════════
#  VLAlert-X scorer (adaptive window, simplified to narrow)
# ═══════════════════════════════════════════════════════════════
class VLAlertXScorer:
    """Score with VLAlertXHead (narrow window only for demo)."""

    def __init__(self, sft_path, x_head_paths, name="VLAlert-X"):
        self.name = name
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.sft_path = sft_path
        self.vlm_loaded = False

        from lkalert.models.components import MultiQueryPMAAggregator
        self.heads = []
        for hp in x_head_paths:
            if not hp.exists():
                continue
            ck = torch.load(hp, weights_only=False, map_location="cpu")
            head_sd = ck["head"]
            d_in = head_sd["aggregator.in_proj.weight"].shape[1]
            head = _build_vlalert_x_head(d_in)
            head.load_state_dict(head_sd)
            head.to(self.device).eval()
            self.heads.append(head)
        logger.info(f"[{name}] {len(self.heads)} VLAlert-X heads loaded")

    def _ensure_vlm(self):
        if self.vlm_loaded:
            return
        logger.info(f"[{self.name}] loading VLM from {self.sft_path}...")
        from transformers import AutoProcessor, AutoModelForImageTextToText
        from peft import PeftModel
        from training.VLA.cot_belief_dataset_v2 import ALL_SPECIAL, BELIEF_OPEN, BELIEF_CLOSE

        self.processor = AutoProcessor.from_pretrained(BASE_MODEL, trust_remote_code=True)
        self.processor.tokenizer.add_special_tokens({"additional_special_tokens": ALL_SPECIAL})
        self.processor.tokenizer.padding_side = "right"
        base = AutoModelForImageTextToText.from_pretrained(
            BASE_MODEL, torch_dtype=torch.bfloat16, trust_remote_code=True)
        base.resize_token_embeddings(len(self.processor.tokenizer))
        self.vlm = PeftModel.from_pretrained(base, self.sft_path).to(self.device)
        self.vlm.eval()
        self.belief_open_id = self.processor.tokenizer.convert_tokens_to_ids(BELIEF_OPEN)
        self.belief_close_id = self.processor.tokenizer.convert_tokens_to_ids(BELIEF_CLOSE)
        self.belief_layers = [20, 24, 28, 32]
        self.vlm_loaded = True
        logger.info(f"[{self.name}] VLM loaded")

    def share_vlm(self, other_scorer):
        """Borrow VLM from another scorer to avoid double-loading."""
        other_scorer._ensure_vlm()
        self.vlm = other_scorer.vlm
        self.processor = other_scorer.processor
        self.belief_open_id = other_scorer.belief_open_id
        self.belief_close_id = other_scorer.belief_close_id
        self.belief_layers = other_scorer.belief_layers
        self.vlm_loaded = True
        self._shared = True
        logger.info(f"[{self.name}] sharing VLM from {other_scorer.name}")

    @torch.no_grad()
    def _extract_belief(self, frames_8):
        self._ensure_vlm()
        from training.VLA.cot_belief_dataset_v2 import SYSTEM_PROMPT_V2, USER_PROMPT_V2
        user_content = [{"type": "image", "image": img} for img in frames_8]
        user_content.append({"type": "text", "text": USER_PROMPT_V2})
        msgs = [
            {"role": "system", "content": [{"type": "text", "text": SYSTEM_PROMPT_V2}]},
            {"role": "user", "content": user_content},
        ]
        text = self.processor.apply_chat_template(msgs, add_generation_prompt=True, tokenize=False)
        inputs = self.processor(text=[text], images=frames_8, return_tensors="pt",
                                padding=True).to(self.device)
        with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
            out = self.vlm(**inputs, output_hidden_states=True, return_dict=True)
        hs_tuple = out.hidden_states
        ids = inputs["input_ids"][0]
        open_pos = (ids == self.belief_open_id).nonzero(as_tuple=False).flatten().tolist()
        close_pos = (ids == self.belief_close_id).nonzero(as_tuple=False).flatten().tolist()
        n_blocks = min(len(open_pos), len(close_pos), 8)
        D = hs_tuple[self.belief_layers[0]].shape[-1]
        belief = torch.zeros(1, 8, len(self.belief_layers) * D, dtype=torch.float16)
        valid = torch.zeros(1, 8, dtype=torch.bool)
        for f in range(n_blocks):
            o, c = open_pos[f], close_pos[f]
            if c <= o + 1:
                continue
            parts = [hs_tuple[L][0, o+1:c].mean(dim=0).to(torch.float16) for L in self.belief_layers]
            belief[0, f] = torch.cat(parts, dim=-1).cpu()
            valid[0, f] = True
        del out, hs_tuple, inputs
        torch.cuda.empty_cache()
        return belief, valid

    @torch.no_grad()
    def score_video(self, video_dir, n_frames, fps, batch_size=2):
        tick_interval = max(1, int(fps))
        tick_frames = list(range(0, n_frames, tick_interval))
        all_frame_sets = []
        for tf in tick_frames:
            end = min(tf + 7, n_frames - 1)
            start = max(0, end - 7)
            indices = list(range(start, end + 1))
            while len(indices) < 8:
                indices = [indices[0]] + indices
            all_frame_sets.append(load_frames(video_dir, indices[:8]))

        results = []
        for bi in tqdm(range(0, len(tick_frames), batch_size),
                       desc=f"{self.name}", ncols=80, leave=False):
            # VLAlert-X scorer: process one at a time (uses same _extract_belief)
            for j in range(min(batch_size, len(tick_frames) - bi)):
                belief, valid = self._extract_belief(all_frame_sets[bi + j])
                b = belief.to(self.device, dtype=torch.float32)
                v = valid.to(self.device)
                probs_all = []
                for head in self.heads:
                    agg_out = head.aggregator(b, v)
                    agg = agg_out[0] if isinstance(agg_out, tuple) else agg_out
                    flat = agg.reshape(1, -1)
                    logits = head.policy_head(flat)
                    probs_all.append(torch.softmax(logits, dim=-1))
                avg = torch.stack(probs_all).mean(dim=0)
                tf = tick_frames[bi + j]
                results.append({"frame": tf, "t": tf / fps,
                                 "p_alert": float(avg[0, 2].cpu()),
                                 "p_observe": float(avg[0, 1].cpu()),
                                 "action": ["SILENT", "OBSERVE", "ALERT"][int(avg.argmax(dim=-1)[0].cpu())]})
        return results

    def unload_vlm(self):
        if self.vlm_loaded and not getattr(self, '_shared', False):
            del self.vlm
            self.vlm_loaded = False
            free_gpu()
            logger.info(f"[{self.name}] VLM unloaded")


def _build_vlalert_x_head(d_in):
    """Build VLAlertXHead architecture from checkpoint dims."""
    from lkalert.models.components import MultiQueryPMAAggregator
    import torch.nn as nn
    K, d_out, hidden = 4, 512, 512
    agg = MultiQueryPMAAggregator(d_in=d_in, d_out=d_out, K=K, n_heads=4)
    policy_head = nn.Sequential(nn.Linear(K * d_out, hidden), nn.GELU(),
                                 nn.Dropout(0.1), nn.Linear(hidden, 3))
    alert_prob_head = nn.Sequential(nn.Linear(K * d_out, hidden // 2), nn.GELU(),
                                     nn.Linear(hidden // 2, 1))
    hazard_head = nn.Linear(K * d_out, 8)
    vjepa_head = nn.Sequential(nn.Linear(K * d_out, hidden), nn.GELU(),
                                nn.Linear(hidden, 1024))
    from lkalert.models.adaptive_window import AdaptiveWindowModule
    wm = AdaptiveWindowModule(belief_dim=d_in)
    head = nn.Module()
    head.aggregator = agg
    head.policy_head = policy_head
    head.alert_prob_head = alert_prob_head
    head.hazard_head = hazard_head
    head.vjepa_head = vjepa_head
    head.window_module = wm
    return head


# ═══════════════════════════════════════════════════════════════
#  M10 scorer (older architecture, single-layer 2560 belief)
# ═══════════════════════════════════════════════════════════════
class M10Scorer:
    """Score with MultiQueryPolicyHead (5-seed ensemble) on B0 backbone."""

    def __init__(self, sft_path, head_paths, name="VLAlert-M10"):
        self.name = name
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.sft_path = sft_path
        self.vlm_loaded = False

        from lkalert.models.components import MultiQueryPolicyHead
        self.heads = []
        for hp in head_paths:
            if not hp.exists():
                continue
            sd = torch.load(hp, weights_only=False, map_location="cpu")
            d_in = sd["aggregator.in_proj.weight"].shape[1]
            head = MultiQueryPolicyHead(hidden_dim=d_in, d_out=512, K=4, n_heads=4)
            head.load_state_dict(sd)
            head.to(self.device).eval()
            self.heads.append(head)
        logger.info(f"[{name}] {len(self.heads)} M10 heads loaded")

    def _ensure_vlm(self):
        if self.vlm_loaded:
            return
        logger.info(f"[{self.name}] loading VLM from {self.sft_path}...")
        from transformers import AutoProcessor, AutoModelForImageTextToText
        from peft import PeftModel
        from training.VLA.cot_belief_dataset_v2 import ALL_SPECIAL

        self.processor = AutoProcessor.from_pretrained(BASE_MODEL, trust_remote_code=True)
        self.processor.tokenizer.add_special_tokens({"additional_special_tokens": ALL_SPECIAL})
        self.processor.tokenizer.padding_side = "right"
        base = AutoModelForImageTextToText.from_pretrained(
            BASE_MODEL, torch_dtype=torch.bfloat16, trust_remote_code=True)
        base.resize_token_embeddings(len(self.processor.tokenizer))
        self.vlm = PeftModel.from_pretrained(base, self.sft_path).to(self.device)
        self.vlm.eval()

        from training.VLA.cot_belief_dataset_v2 import BELIEF_OPEN, BELIEF_CLOSE
        tok = self.processor.tokenizer
        self.action_ids = set()
        for t in ["<|ACTION_SILENT|>", "<|ACTION_OBSERVE|>", "<|ACTION_ALERT|>"]:
            tid = tok.convert_tokens_to_ids(t)
            if tid != tok.unk_token_id:
                self.action_ids.add(tid)
        self.belief_open_id = tok.convert_tokens_to_ids(BELIEF_OPEN)
        self.belief_close_id = tok.convert_tokens_to_ids(BELIEF_CLOSE)
        self.vlm_loaded = True
        logger.info(f"[{self.name}] VLM loaded (single-layer 2560 extraction)")

    @torch.no_grad()
    def _extract_belief(self, frames_8):
        """Extract last-layer belief [1, 8, 2560] using action-token positions."""
        self._ensure_vlm()
        from training.VLA.cot_belief_dataset_v2 import SYSTEM_PROMPT_V2, USER_PROMPT_V2
        user_content = [{"type": "image", "image": img} for img in frames_8]
        user_content.append({"type": "text", "text": USER_PROMPT_V2})
        msgs = [
            {"role": "system", "content": [{"type": "text", "text": SYSTEM_PROMPT_V2}]},
            {"role": "user", "content": user_content},
        ]
        text = self.processor.apply_chat_template(msgs, add_generation_prompt=True, tokenize=False)
        inputs = self.processor(text=[text], images=frames_8, return_tensors="pt",
                                padding=True).to(self.device)
        with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
            out = self.vlm(**inputs, output_hidden_states=True, return_dict=True)
        hs_last = out.hidden_states[-1][0]  # [T, 2560]
        ids = inputs["input_ids"][0]

        action_pos = [int(p) for p, t in enumerate(ids.tolist()) if t in self.action_ids]
        if len(action_pos) < 1:
            close_pos = (ids == self.belief_close_id).nonzero(as_tuple=False).flatten().tolist()
            action_pos = close_pos

        D = hs_last.shape[-1]
        belief = torch.zeros(1, 8, D, dtype=torch.float16)
        valid = torch.zeros(1, 8, dtype=torch.bool)
        for f in range(min(len(action_pos), 8)):
            belief[0, f] = hs_last[action_pos[f]].to(torch.float16).cpu()
            valid[0, f] = True
        del out, inputs, hs_last
        torch.cuda.empty_cache()
        return belief, valid

    @torch.no_grad()
    def score_video(self, video_dir, n_frames, fps, batch_size=2):
        tick_interval = max(1, int(fps))
        tick_frames = list(range(0, n_frames, tick_interval))
        all_frame_sets = []
        for tf in tick_frames:
            end = min(tf + 7, n_frames - 1)
            start = max(0, end - 7)
            indices = list(range(start, end + 1))
            while len(indices) < 8:
                indices = [indices[0]] + indices
            all_frame_sets.append(load_frames(video_dir, indices[:8]))

        results = []
        prev_action = torch.tensor([0], device=self.device, dtype=torch.long)
        for bi in tqdm(range(0, len(tick_frames)),
                       desc=f"{self.name}", ncols=80, leave=False):
            belief, valid = self._extract_belief(all_frame_sets[bi])
            b = belief.to(self.device, dtype=torch.float32)
            v = valid.to(self.device)
            tta_m = torch.tensor([5.0], device=self.device)
            tta_v = torch.tensor([1.0], device=self.device)

            probs_all = []
            for head in self.heads:
                logits, _ = head(b, v, tta_m, tta_v, prev_action)
                probs_all.append(torch.softmax(logits, dim=-1))

            avg = torch.stack(probs_all).mean(dim=0)
            p_alert = float(avg[0, 2].cpu())
            p_obs = float(avg[0, 1].cpu())
            action_idx = int(avg.argmax(dim=-1)[0].cpu())
            action = ["SILENT", "OBSERVE", "ALERT"][action_idx]
            prev_action = torch.tensor([action_idx], device=self.device, dtype=torch.long)
            tf = tick_frames[bi]
            results.append({"frame": tf, "t": tf / fps,
                             "p_alert": p_alert, "p_observe": p_obs, "action": action})
        return results

    def unload_vlm(self):
        if self.vlm_loaded:
            del self.vlm
            self.vlm_loaded = False
            free_gpu()
            logger.info(f"[{self.name}] VLM unloaded")


# ═══════════════════════════════════════════════════════════════
#  Qwen2.5-VL-3B scorer (monolithic TTA head)
# ═══════════════════════════════════════════════════════════════
class Qwen25Scorer:
    """Score with Qwen2.5-VL-3B + TTAHead (TTA regression → threshold → action)."""

    def __init__(self, name="VLAlert-2.5"):
        self.name = name
        self.device = "cuda"
        self.vlm = None

    def _load(self):
        if self.vlm is not None:
            return
        logger.info(f"[{self.name}] loading Qwen2.5-VL-3B...")
        from transformers import AutoProcessor, AutoModelForImageTextToText
        from peft import PeftModel
        import torch.nn as nn
        import torch.nn.functional as F

        self.processor = AutoProcessor.from_pretrained(
            BASE_MODEL_Q25, trust_remote_code=True)
        self.processor.tokenizer.padding_side = "right"

        base = AutoModelForImageTextToText.from_pretrained(
            BASE_MODEL_Q25, torch_dtype=torch.bfloat16, trust_remote_code=True)
        self.vlm = PeftModel.from_pretrained(base, SFT_Q25_LORA).to(self.device)
        self.vlm.eval()

        class TTAHead(nn.Module):
            def __init__(self, hidden_dim, intermediate_dim=512):
                super().__init__()
                self.net = nn.Sequential(
                    nn.Linear(hidden_dim, intermediate_dim), nn.GELU(), nn.Dropout(0.1),
                    nn.Linear(intermediate_dim, intermediate_dim // 2), nn.GELU(), nn.Dropout(0.1),
                    nn.Linear(intermediate_dim // 2, 2),
                )
            def forward(self, h):
                out = self.net(h)
                return F.softplus(out[:, 0]), out[:, 1]

        self.tta_head = TTAHead(2048, 512).to(self.device)
        sd = torch.load(TTA_HEAD_Q25, weights_only=False, map_location="cpu")
        self.tta_head.load_state_dict(sd)
        self.tta_head.eval()
        logger.info(f"[{self.name}] loaded, GPU: {torch.cuda.memory_allocated()//1024**2}MB")

    @torch.no_grad()
    def _score_batch(self, frame_sets):
        self._load()
        N = len(frame_sets)
        texts, all_images = [], []
        for frames_8 in frame_sets:
            uc = [{"type": "image", "image": img} for img in frames_8]
            uc.append({"type": "text", "text": "Describe the driving safety situation."})
            msgs = [{"role": "user", "content": uc}]
            texts.append(self.processor.apply_chat_template(
                msgs, add_generation_prompt=True, tokenize=False))
            all_images.extend(frames_8)

        inputs = self.processor(text=texts, images=all_images,
                                return_tensors="pt", padding=True).to(self.device)

        core = self.vlm.get_base_model().model
        with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
            out = core(
                input_ids=inputs["input_ids"],
                attention_mask=inputs.get("attention_mask"),
                pixel_values=inputs.get("pixel_values"),
                image_grid_thw=inputs.get("image_grid_thw"),
                use_cache=False, return_dict=True,
            )
        hs = out.last_hidden_state  # [N, L, 2048]
        mask = inputs["attention_mask"].unsqueeze(-1).to(hs.dtype)
        belief = (hs * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1)  # [N, 2048]
        tta_mean, _ = self.tta_head(belief.float())  # [N]

        results = []
        for i in range(N):
            tta = float(tta_mean[i].cpu())
            if tta < 2.0:
                action = "ALERT"
            elif tta < 5.0:
                action = "OBSERVE"
            else:
                action = "SILENT"
            p_alert = max(0.0, min(1.0, 1.0 - tta / 10.0))
            results.append((p_alert, action, tta))
        return results

    def score_video(self, video_dir, n_frames, fps, batch_size=2):
        tick_interval = max(1, int(fps))
        tick_frames = list(range(0, n_frames, tick_interval))
        all_frame_sets = []
        for tf in tick_frames:
            end = min(tf + 7, n_frames - 1)
            start = max(0, end - 7)
            indices = list(range(start, end + 1))
            while len(indices) < 8:
                indices = [indices[0]] + indices
            all_frame_sets.append(load_frames(video_dir, indices[:8]))

        results = []
        for bi in tqdm(range(0, len(tick_frames), batch_size),
                       desc=f"{self.name}", ncols=80, leave=False):
            batch = all_frame_sets[bi:bi + batch_size]
            batch_results = self._score_batch(batch)
            for j, (p_alert, action, tta) in enumerate(batch_results):
                tf = tick_frames[bi + j]
                results.append({"frame": tf, "t": tf / fps,
                                 "p_alert": p_alert, "action": action,
                                 "tta_mean": tta})
        return results

    def unload_vlm(self):
        if self.vlm is not None:
            del self.vlm, self.tta_head
            self.vlm = None
            free_gpu()
            logger.info(f"[{self.name}] unloaded")


# ═══════════════════════════════════════════════════════════════
#  Visualization
# ═══════════════════════════════════════════════════════════════
ACTION_COLORS = {"SILENT": (0, 200, 0), "OBSERVE": (0, 200, 255), "ALERT": (0, 0, 255)}

def render_comparison_video(video_dir: Path, model_scores: dict[str, list[dict]],
                              fps: float, n_frames: int, out_path: Path):
    """Render a comparison video: left=frame, right=score curves + actions."""
    W_FRAME = 640
    H_FRAME = 360
    W_PANEL = 400
    W_TOTAL = W_FRAME + W_PANEL
    H_TOTAL = H_FRAME

    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
    writer = cv2.VideoWriter(str(out_path), fourcc, min(fps, 30), (W_TOTAL, H_TOTAL))

    # Precompute score arrays interpolated to native fps
    model_names = list(model_scores.keys())
    colors_bgr = [
        (255, 100, 100),   # blue-ish for BADAS
        (100, 255, 100),   # green for VLAlert-v3
        (0, 180, 255),     # orange for VLAlert-v2
        (100, 100, 255),   # red for VLAlert-X
        (255, 255, 100),   # cyan for VLAlert-M10
        (200, 100, 255),   # pink
    ]

    # Interpolate each model's p_alert to native fps
    interp_scores = {}
    interp_actions = {}
    for mname, results in model_scores.items():
        if not results: continue
        tick_frames = [r["frame"] for r in results]
        tick_palert = [r["p_alert"] for r in results]
        tick_actions = [r["action"] for r in results]
        # Interpolate p_alert to every frame
        all_p = np.interp(range(n_frames), tick_frames, tick_palert)
        interp_scores[mname] = all_p
        # Nearest-neighbor for actions
        all_a = []
        for f in range(n_frames):
            closest = min(range(len(tick_frames)), key=lambda i: abs(tick_frames[i] - f))
            all_a.append(tick_actions[closest])
        interp_actions[mname] = all_a

    # History window for score plot (last 5 seconds)
    history_frames = int(5 * fps)

    for f in tqdm(range(n_frames), desc="render", ncols=80, leave=False):
        # Load frame
        frame_path = video_dir / f"{f:06d}.jpg"
        if frame_path.exists():
            img = cv2.imread(str(frame_path))
            img = cv2.resize(img, (W_FRAME, H_FRAME))
        else:
            img = np.zeros((H_FRAME, W_FRAME, 3), dtype=np.uint8)

        # Create right panel (white background)
        panel = np.ones((H_TOTAL, W_PANEL, 3), dtype=np.uint8) * 240

        # Draw score curves
        t_sec = f / fps
        plot_y0 = 30
        plot_y1 = H_TOTAL - 80
        plot_h = plot_y1 - plot_y0
        plot_x0 = 10
        plot_x1 = W_PANEL - 10
        plot_w = plot_x1 - plot_x0

        # Grid lines
        for y_val in [0.0, 0.25, 0.5, 0.75, 1.0]:
            y = int(plot_y1 - y_val * plot_h)
            cv2.line(panel, (plot_x0, y), (plot_x1, y), (200, 200, 200), 1)
            cv2.putText(panel, f"{y_val:.1f}", (plot_x1 + 2, y + 4),
                         cv2.FONT_HERSHEY_SIMPLEX, 0.3, (128, 128, 128), 1)

        # Title
        cv2.putText(panel, f"t={t_sec:.1f}s", (plot_x0, 20),
                     cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1)

        # Draw each model's score curve
        win_start = max(0, f - history_frames)
        for mi, mname in enumerate(model_names):
            if mname not in interp_scores: continue
            scores = interp_scores[mname]
            color = colors_bgr[mi % len(colors_bgr)]

            # Draw curve
            for x in range(plot_w - 1):
                fi = win_start + int(x * (f - win_start + 1) / plot_w)
                fi_next = win_start + int((x + 1) * (f - win_start + 1) / plot_w)
                fi = min(fi, n_frames - 1)
                fi_next = min(fi_next, n_frames - 1)
                y1 = int(plot_y1 - scores[fi] * plot_h)
                y2 = int(plot_y1 - scores[fi_next] * plot_h)
                cv2.line(panel, (plot_x0 + x, y1), (plot_x0 + x + 1, y2), color, 2)

            # Current action label
            action = interp_actions[mname][f] if mname in interp_actions else "?"
            label_y = H_TOTAL - 70 + mi * 18
            act_color = ACTION_COLORS.get(action, (128, 128, 128))
            cv2.putText(panel, f"{mname}: ", (5, label_y),
                         cv2.FONT_HERSHEY_SIMPLEX, 0.4, (0, 0, 0), 1)
            cv2.putText(panel, f"{action} ({scores[f]:.2f})", (5 + len(mname) * 8, label_y),
                         cv2.FONT_HERSHEY_SIMPLEX, 0.4, act_color[::-1], 1)

        # Combine frame + panel
        combined = np.hstack([img, panel])
        writer.write(combined)

    writer.release()
    logger.info(f"  saved → {out_path}")


# ═══════════════════════════════════════════════════════════════
#  Main
# ═══════════════════════════════════════════════════════════════
def get_video_info(video_dir: Path):
    frames = sorted(video_dir.glob("*.jpg"))
    n = len(frames)
    # Try to detect fps from parent video
    parent_video = None
    for ext in [".mp4", ".avi"]:
        p = ROOT / "demo/compare" / (video_dir.name + ext)
        if p.exists(): parent_video = p; break
    fps = 30.0
    if parent_video:
        cap = cv2.VideoCapture(str(parent_video))
        fps = cap.get(cv2.CAP_PROP_FPS) or 30.0
        cap.release()
    return n, fps


def score_one_model(mname, scorer, videos, batch_size=2):
    """Score all videos with one model, save incrementally."""
    total_ticks = 0
    t0_all = time.time()
    for video_dir in videos:
        vname = video_dir.name
        n_frames, fps = get_video_info(video_dir)
        scores_path = OUT_DIR / vname / "scores.json"
        scores_path.parent.mkdir(parents=True, exist_ok=True)
        cached = json.loads(scores_path.read_text()) if scores_path.exists() else {}
        if mname in cached:
            logger.info(f"  [{mname}] {vname}: cached ({len(cached[mname])} ticks)")
            total_ticks += len(cached[mname])
            continue
        logger.info(f"  [{mname}] {vname}: {n_frames} frames @ {fps:.0f}fps...")
        t0 = time.time()
        results = scorer.score_video(video_dir, n_frames, fps, batch_size=batch_size)
        dt = time.time() - t0
        cached[mname] = results
        scores_path.write_text(json.dumps(cached, indent=2))
        total_ticks += len(results)
        logger.info(f"  [{mname}] {vname}: {len(results)} ticks in {dt:.1f}s")
    dt_all = time.time() - t0_all
    logger.info(f"  [{mname}] done — {total_ticks} ticks total in {dt_all:.1f}s")


def render_all_videos(videos, model_names):
    """Re-render comparison videos using all cached scores."""
    for video_dir in videos:
        vname = video_dir.name
        n_frames, fps = get_video_info(video_dir)
        scores_path = OUT_DIR / vname / "scores.json"
        if not scores_path.exists():
            continue
        cached = json.loads(scores_path.read_text())
        all_scores = {m: cached[m] for m in model_names if m in cached}
        if not all_scores:
            continue
        any_alert = any(
            any(r["action"] in ("ALERT", "OBSERVE") for r in results)
            for results in all_scores.values()
        )
        if not any_alert:
            logger.info(f"  {vname}: all SILENT, skip viz")
            continue
        out_video = OUT_DIR / vname / "comparison.mp4"
        logger.info(f"  {vname}: rendering with {list(all_scores.keys())}...")
        render_comparison_video(video_dir, all_scores, fps, n_frames, out_video)


ALL_MODELS = ["BADAS", "VLAlert-v3", "VLAlert-v2", "VLAlert-X", "VLAlert-M10"]


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--models", type=str, default="v3,v2,X,M10,q25",
                    help="comma-separated: BADAS,v3,v2,X,M10,q25")
    ap.add_argument("--only", type=str, default="", help="process only this video name")
    ap.add_argument("--batch_size", type=int, default=2,
                    help="VLM batch size (2 fills ~28GB on 32GB GPU)")
    ap.add_argument("--skip_render", action="store_true")
    args = ap.parse_args()

    videos = sorted([d for d in FRAMES_DIR.iterdir() if d.is_dir()])
    if args.only:
        videos = [v for v in videos if args.only in v.name]
    logger.info(f"Processing {len(videos)} videos")

    model_sel = set(args.models.split(","))
    scored_names = []

    # ── Group 0: BADAS (V-JEPA, separate backbone) ──
    if "BADAS" in model_sel:
        logger.info("\n" + "=" * 60 + "\n  BADAS (V-JEPA2)\n" + "=" * 60)
        scorer = BADASScorer()
        score_one_model("BADAS", scorer, videos, batch_size=1)
        scored_names.append("BADAS")
        del scorer
        free_gpu()

    # ── Group 1: VLAlert-v3 (B3 backbone: sft_x_v3) ──
    if "v3" in model_sel:
        logger.info("\n" + "=" * 60 + "\n  VLAlert-v3 (B3: sft_x_v3)\n" + "=" * 60)
        scorer = VLAlertScorer(sft_path=SFT_V3, danger_path=DANGER_V3,
                                policy_paths=[POLICY_V3], name="VLAlert-v3")
        score_one_model("VLAlert-v3", scorer, videos, batch_size=args.batch_size)
        scored_names.append("VLAlert-v3")
        scorer.unload_vlm()
        del scorer
        free_gpu()

    # ── Group 2: VLAlert-v2 + VLAlert-X (B2 backbone: sft_x_v2, shared VLM) ──
    run_v2 = "v2" in model_sel
    run_x = "X" in model_sel
    if run_v2 or run_x:
        logger.info("\n" + "=" * 60 + "\n  B2 backbone group (sft_x_v2)\n" + "=" * 60)
        v2_scorer = None
        x_scorer = None
        if run_v2:
            v2_paths = [p for p in POLICY_V2_SEEDS if p.exists()]
            if v2_paths:
                v2_scorer = VLAlertScorer(sft_path=SFT_V2, danger_path=DANGER_V2,
                                           policy_paths=v2_paths, name="VLAlert-v2")
        if run_x:
            x_paths = [p for p in POLICY_X_SEEDS if p.exists()]
            if x_paths:
                x_scorer = VLAlertXScorer(sft_path=SFT_V2, x_head_paths=x_paths,
                                           name="VLAlert-X")

        # Score VLAlert-v2 first (loads B2 VLM)
        if v2_scorer:
            score_one_model("VLAlert-v2", v2_scorer, videos, batch_size=args.batch_size)
            scored_names.append("VLAlert-v2")

        # Score VLAlert-X sharing B2 VLM from v2
        if x_scorer:
            if v2_scorer and v2_scorer.vlm_loaded:
                x_scorer.share_vlm(v2_scorer)
            score_one_model("VLAlert-X", x_scorer, videos, batch_size=args.batch_size)
            scored_names.append("VLAlert-X")

        if v2_scorer:
            v2_scorer.unload_vlm()
            del v2_scorer
        if x_scorer:
            del x_scorer
        free_gpu()

    # ── Group 3: VLAlert-M10 (B0 backbone: qwen3vl4b_cot_belief_perframe) ──
    if "M10" in model_sel:
        logger.info("\n" + "=" * 60 + "\n  VLAlert-M10 (B0: perframe)\n" + "=" * 60)
        m10_paths = [p for p in M10_SEEDS if p.exists()]
        if m10_paths:
            scorer = M10Scorer(sft_path=SFT_B0, head_paths=m10_paths, name="VLAlert-M10")
            score_one_model("VLAlert-M10", scorer, videos, batch_size=args.batch_size)
            scored_names.append("VLAlert-M10")
            scorer.unload_vlm()
            del scorer
            free_gpu()

    # ── Group 4: VLAlert-2.5 (Qwen2.5-VL-3B, monolithic TTA) ──
    if "q25" in model_sel:
        logger.info("\n" + "=" * 60 + "\n  VLAlert-2.5 (Qwen2.5-VL-3B)\n" + "=" * 60)
        scorer = Qwen25Scorer(name="VLAlert-2.5")
        score_one_model("VLAlert-2.5", scorer, videos, batch_size=args.batch_size)
        scored_names.append("VLAlert-2.5")
        scorer.unload_vlm()
        del scorer
        free_gpu()

    # ── Render comparison videos with all scored models ──
    if not args.skip_render:
        # Include previously cached BADAS too
        render_names = ["BADAS"] + scored_names if "BADAS" not in scored_names else scored_names
        logger.info(f"\n{'='*60}\n  Rendering comparisons: {render_names}\n{'='*60}")
        render_all_videos(videos, render_names)

    logger.info(f"\n✅ All done! Results in {OUT_DIR}")


if __name__ == "__main__":
    main()