training/Policy/multichannel_dataset.py

"""Multi-channel dataset for LKAlert-MCB.

**Design (post Day-10 pivot 2026-04-27):**
LKAlert-MCB is a 2-channel architecture for the headline:
  - Channel 1 (Qwen semantic):  belief_frame [B,T,2560] + valid + text
  - Channel 3 (V-JEPA dynamics): clip-level vjepa_feature [B,1024] (mean
                                  pooled from per-frame [16,1024])

Channel 2 (object motion) is intentionally NOT a learned input here —
it failed Red Line 4 gate on Day 10. Object features remain on disk
for taxonomy / qualitative figures / appendix Table 6.

**Per-cache joins are by clip_id, never by index** (Rule 3).

Each row returned by the dataset is a dict:
  {belief, valid, text, vjepa, vjepa_mask, tta_mean, tta_var, vid, y_p_any}

When V-JEPA features are missing for a clip, `vjepa = zeros(1024)` and
`vjepa_mask = 0` so the fusion MLP can learn to ignore the channel.
"""
from __future__ import annotations

import json
import logging
from pathlib import Path
from typing import Dict, List, Optional

import numpy as np
import torch
from torch.utils.data import Dataset

logger = logging.getLogger("multichannel_dataset")

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


def _diag_filename(cache_name: str) -> str:
    if cache_name.endswith("_diag"):
        return f"{cache_name}.json"
    return f"{cache_name}_diag.json"


def _resolve_vjepa_short(vid: str) -> str:
    """V-JEPA dicts use the short Nexar id (no `nexar_` prefix)."""
    return vid.replace("nexar_", "")


class MultichannelDataset(Dataset):
    """Joins Qwen belief cache + V-JEPA features (and optional labels)."""

    def __init__(self,
                  cache_name: str,
                  split: str,
                  vjepa_path: Optional[Path] = None,
                  with_labels: bool = True):
        self.cache_name = cache_name
        self.split = split

        cache_path = CACHE_DIR / f"{cache_name}.pt"
        if not cache_path.exists():
            raise FileNotFoundError(f"Qwen cache not found: {cache_path}")
        c = torch.load(cache_path, weights_only=False, map_location="cpu")

        self.bf  = c["beliefs_frame"].float()       # [N, T, D]
        self.vf  = c["valid_frames"].bool()          # [N, T]
        self.bt  = c["beliefs_text"].float()         # [N, D]
        self.tm  = c["tta_means"].float()            # [N]
        self.tv  = c["tta_vars"].float()             # [N]
        self.ids = c["meta"]["ids"]
        self.action_labels = c["meta"].get("action_labels", [])

        # ── V-JEPA feature dict ──────────────────────────────────────────
        if vjepa_path is None:
            # default heuristic: train→multisrc train, val→multisrc val,
            # test→clip features
            if split == "train":
                vjepa_path = VJEPA_DIR / "train_perframe_multisrc.pt"
            elif split == "val":
                vjepa_path = VJEPA_DIR / "val_perframe_multisrc.pt"
            else:
                vjepa_path = VJEPA_DIR / "test_clip_features.pt"

        if not vjepa_path.exists():
            logger.warning(f"  V-JEPA cache {vjepa_path} missing — "
                            "all V-JEPA features will be zero-masked")
            self.vj_dict = {}
        else:
            self.vj_dict = torch.load(vjepa_path, weights_only=False,
                                          map_location="cpu")

        # detect per-frame [T, 1024] vs clip-level [1024]
        if self.vj_dict:
            sample = next(iter(self.vj_dict.values()))
            self.vj_per_frame = (sample.dim() == 2)
        else:
            self.vj_per_frame = False

        # pre-compute clip-level V-JEPA per id (mean-pool if per-frame)
        self.vj_clip: Dict[str, torch.Tensor] = {}
        for vid in self.ids:
            short = _resolve_vjepa_short(vid)
            v = self.vj_dict.get(short)
            if v is None:
                continue
            v = v.float()
            if v.dim() == 2:
                v = v.mean(dim=0)
            self.vj_clip[vid] = v
        cov = len(self.vj_clip) / max(1, len(self.ids))
        logger.info(f"[mcb-dataset:{cache_name}] N={len(self.ids)}  "
                    f"V-JEPA coverage={100*cov:.1f}%  "
                    f"vj_per_frame={self.vj_per_frame}")

        # ── Labels (optional) ────────────────────────────────────────────
        if with_labels:
            diag_path = DIAG_DIR / _diag_filename(cache_name)
            if diag_path.exists():
                raw = json.loads(diag_path.read_text())
                by_id = {s["video_id"]: s for s in raw["samples"]}
                self.y_any = np.asarray(
                    [1 if by_id.get(v, {}).get("action_label") == 2 else 0
                      for v in self.ids], dtype=np.float32)
            else:
                # fallback: from cache action_labels
                self.y_any = np.asarray(
                    [1 if a == 2 else 0 for a in self.action_labels],
                    dtype=np.float32) if self.action_labels else np.zeros(
                    len(self.ids), dtype=np.float32)
        else:
            self.y_any = None

    def __len__(self) -> int:
        return len(self.ids)

    def __getitem__(self, i: int) -> Dict:
        vid = self.ids[i]
        vj = self.vj_clip.get(vid)
        if vj is None:
            vj = torch.zeros(1024, dtype=torch.float32)
            mask = torch.tensor(0.0)
        else:
            mask = torch.tensor(1.0)
        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),
            "vjepa":    vj,
            "vjepa_mask": mask,
            "vid":      vid,
        }
        if self.y_any is not None:
            out["y_p_any"] = torch.tensor(float(self.y_any[i]),
                                            dtype=torch.float32)
        return out


def collate(batch: List[Dict]) -> Dict:
    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]),
        "vjepa":    torch.stack([b["vjepa"]    for b in batch]),
        "vjepa_mask": torch.stack([b["vjepa_mask"] for b in batch]),
        "vids":     [b["vid"] for b in batch],
    }
    if "y_p_any" in batch[0]:
        out["y_p_any"] = torch.stack([b["y_p_any"] for b in batch])
    return out