tools/build_unified_benchmark.py

"""Build VLAlert-Bench unified benchmark.

Pipeline:
  Step 1: scan 6 source datasets -> per-video splits
  Step 2: per-frame action labels per (positive) video
  Step 3: 1Hz tick-level parquet (train/val/test/extra_val_adasto/extra_val_accident)
  Step 4: HF dataset card README.md + loader vlalert_bench.py
  Step 5: leakage verification + smoke test

Usage:
    python tools/build_unified_benchmark.py --step 1        # video splits only
    python tools/build_unified_benchmark.py --step 2        # add frame labels
    python tools/build_unified_benchmark.py --step 3        # add tick parquet
    python tools/build_unified_benchmark.py --step 4        # HF card + loader
    python tools/build_unified_benchmark.py --step 5        # verify
    python tools/build_unified_benchmark.py --step all      # do everything
"""
from __future__ import annotations
import argparse
import json
import logging
import random
from collections import Counter, defaultdict
from pathlib import Path
from typing import Dict, List, Optional, Tuple

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

# ───────────────────────────── paths ─────────────────────────────
ROOT = Path("PROJECT_ROOT")
NEXAR_DIR = ROOT / "NEXAR_COLLISION"
DAD_DIR = ROOT / "DAD" / "videos"
DOTA_DIR = ROOT / "DoTA"
DADA_DIR = ROOT / "DADA-2000"
ADASTO_DIR = ROOT / "ADAS-TO-Critic"
CARLA_DIR = ROOT / "accident"

BENCH_DIR = ROOT / "benchmark" / "v1"
MANIFEST_DIR = BENCH_DIR / "manifest"
DATA_DIR = BENCH_DIR / "data"
STATS_DIR = BENCH_DIR / "stats"

# Reproducibility
SEED = 42

# ───────────────────── Step 1: video splits ─────────────────────


def collect_nexar() -> Dict[str, Dict]:
    """Returns video_id -> {split, category, source_dir, source} for Nexar."""
    out = {}
    split_map = {
        "train": "train",
        "test-public": "val",       # → in-domain VAL
        "test-private": "test",     # → in-domain TEST
    }
    cat_map = {"positive": "ego_positive", "negative": "safe_neg"}
    for src_split, dst_split in split_map.items():
        for cat_dir, cat_label in cat_map.items():
            d = NEXAR_DIR / src_split / cat_dir
            if not d.exists():
                continue
            for vid_path in sorted(d.glob("*.mp4")):
                vid_id = f"nexar_{vid_path.stem}"
                out[vid_id] = {
                    "video_id": vid_id,
                    "source": "nexar",
                    "split": dst_split,
                    "category": cat_label,
                    "video_path": str(vid_path.relative_to(ROOT)),
                    "native_split": src_split,
                }
    return out


def collect_dad(seed: int = SEED, val_frac: float = 0.10) -> Dict[str, Dict]:
    """DAD: native training -> 90% train + 10% val (stratified by category);
    native testing -> test."""
    out = {}
    cat_map = {"positive": "ego_positive", "negative": "safe_neg"}
    # 1. testing -> test (untouched)
    for cat_dir, cat_label in cat_map.items():
        d = DAD_DIR / "testing" / cat_dir
        if not d.exists():
            continue
        for vid_path in sorted(d.glob("*.mp4")):
            vid_id = f"dad_testi_{cat_dir[:3]}_{vid_path.stem}"
            out[vid_id] = {
                "video_id": vid_id,
                "source": "dad",
                "split": "test",
                "category": cat_label,
                "video_path": str(vid_path.relative_to(ROOT)),
                "native_split": "testing",
            }
    # 2. training -> 90% train + 10% val, stratified
    for cat_dir, cat_label in cat_map.items():
        d = DAD_DIR / "training" / cat_dir
        if not d.exists():
            continue
        vids = sorted(d.glob("*.mp4"))
        rng = random.Random(seed + hash(("dad", cat_label)) % 1000)
        ids = [p.stem for p in vids]
        rng.shuffle(ids)
        n_val = max(1, int(len(ids) * val_frac))
        val_set = set(ids[:n_val])
        for vid_path in vids:
            stem = vid_path.stem
            vid_id = f"dad_train_{cat_dir[:3]}_{stem}"
            out[vid_id] = {
                "video_id": vid_id,
                "source": "dad",
                "split": "val" if stem in val_set else "train",
                "category": cat_label,
                "video_path": str(vid_path.relative_to(ROOT)),
                "native_split": "training",
            }
    return out


def collect_dota(seed: int = SEED, val_frac: float = 0.10) -> Dict[str, Dict]:
    """DoTA: metadata_train -> 90% train + 10% val (stratified ego/non-ego);
    metadata_val -> test (held out, untouched)."""
    out = {}
    # 1. metadata_val -> test (untouched)
    val_meta = DOTA_DIR / "metadata_val.json"
    if val_meta.exists():
        meta = json.load(open(val_meta))
        for k, v in meta.items():
            ego = "ego" in v.get("anomaly_class", "").lower()
            cat = "ego_positive" if ego else "non_ego"
            out[f"dota_{k}"] = {
                "video_id": f"dota_{k}",
                "source": "dota",
                "split": "test",
                "category": cat,
                "video_path": str((DOTA_DIR / "frames" / k).relative_to(ROOT)),
                "anomaly_class": v.get("anomaly_class"),
                "anomaly_start": v.get("anomaly_start"),
                "anomaly_end": v.get("anomaly_end"),
                "num_frames": v.get("num_frames"),
                "native_split": "metadata_val",
            }
    # 2. metadata_train -> 90% train + 10% val, stratified by category
    train_meta = DOTA_DIR / "metadata_train.json"
    if train_meta.exists():
        meta = json.load(open(train_meta))
        # bucket by category for stratified split
        buckets: Dict[str, List[str]] = defaultdict(list)
        for k, v in meta.items():
            ego = "ego" in v.get("anomaly_class", "").lower()
            cat = "ego_positive" if ego else "non_ego"
            buckets[cat].append(k)
        val_set = set()
        for cat, keys in buckets.items():
            rng = random.Random(seed + hash(("dota", cat)) % 1000)
            keys_shuf = list(keys)
            rng.shuffle(keys_shuf)
            n_val = max(1, int(len(keys_shuf) * val_frac))
            val_set.update(keys_shuf[:n_val])
        for k, v in meta.items():
            ego = "ego" in v.get("anomaly_class", "").lower()
            cat = "ego_positive" if ego else "non_ego"
            out[f"dota_{k}"] = {
                "video_id": f"dota_{k}",
                "source": "dota",
                "split": "val" if k in val_set else "train",
                "category": cat,
                "video_path": str((DOTA_DIR / "frames" / k).relative_to(ROOT)),
                "anomaly_class": v.get("anomaly_class"),
                "anomaly_start": v.get("anomaly_start"),
                "anomaly_end": v.get("anomaly_end"),
                "num_frames": v.get("num_frames"),
                "native_split": "metadata_train",
            }
    return out


def collect_dada(seed: int = SEED) -> Dict[str, Dict]:
    """DADA-2000: random 80/10/10 by video_id (positive + negative); non-ego excluded.

    Per-video annotation.json is loaded later in Step 2; here we only need
    the split assignment.
    """
    out = {}
    cat_dirs = {
        "positive": "ego_positive",
        "negative": "safe_neg",
        "non-ego":  "non_ego",
    }
    # group video_ids by category for stratified split
    for cat_dir, cat_label in cat_dirs.items():
        d = DADA_DIR / cat_dir
        if not d.exists():
            continue
        # each video is a folder like images_10_001/
        vid_dirs = sorted([p for p in d.iterdir() if p.is_dir()])
        vid_ids = [p.name for p in vid_dirs]
        rng = random.Random(seed + hash(cat_label) % 1000)
        rng.shuffle(vid_ids)
        n = len(vid_ids)
        n_train = int(n * 0.80)
        n_val = int(n * 0.10)
        # non-ego: still gets a split but flagged as excluded from main pool
        for i, vid_name in enumerate(vid_ids):
            if i < n_train:
                dst = "train"
            elif i < n_train + n_val:
                dst = "val"
            else:
                dst = "test"
            vid_id = f"dada_{vid_name}"
            out[vid_id] = {
                "video_id": vid_id,
                "source": "dada",
                "split": dst,
                "category": cat_label,
                "video_path": str((DADA_DIR / cat_dir / vid_name).relative_to(ROOT)),
                "native_split": None,
                "excluded_from_main": (cat_label == "non_ego"),
            }
    return out


def collect_adasto() -> Dict[str, Dict]:
    """ADAS-TO-Critic: all videos go to extra_val_adasto (held-out OOD).

    All clips are uniformly 20 s with takeover at t = 10 s; we expose the
    entire corpus as a single held-out OOD split — it is never used for
    training or model selection."""
    out = {}
    for vid_path in sorted(ADASTO_DIR.glob("*.mp4")):
        vid_name = vid_path.stem
        vid_id = f"adasto_{vid_name}"
        out[vid_id] = {
            "video_id": vid_id,
            "source": "adasto_critic",
            "split": "extra_val_adasto",
            "category": "mixed",
            "video_path": str(vid_path.relative_to(ROOT)),
            "native_split": None,
            "t_takeover_s": 10.0,
            "duration_s": 20.0,
        }
    return out


def collect_accident() -> Dict[str, Dict]:
    """Kaggle ACCIDENT @ CVPR 2026 (Picek et al.) -> extra_val_accident only.

    Source: https://www.kaggle.com/competitions/accident
    Clips are rendered with CARLA but are released under the Kaggle ACCIDENT
    competition by Picek et al.; we treat them as a held-out OOD test set."""
    import csv
    out = {}
    manifest_csv = CARLA_DIR / "takeover_manifest.csv"
    if not manifest_csv.exists():
        logger.warning(f"ACCIDENT manifest not found: {manifest_csv}")
        return out
    with manifest_csv.open() as f:
        for row in csv.DictReader(f):
            clip = row.get("clip", "").strip()
            if not clip:
                continue
            vid_id = f"accident_{clip}"
            out[vid_id] = {
                "video_id": vid_id,
                "source": "accident",
                "split": "extra_val_accident",
                "category": "ego_positive",
                "video_path": str((CARLA_DIR / "sim_dataset" / "videos" /
                                    row.get("accident_type", "") / f"{clip}.mp4").relative_to(ROOT)),
                "native_split": None,
                "t_takeover_s": float(row.get("t_takeover", 0)),
                "accident_type": row.get("accident_type"),
                "weather": row.get("weather"),
                "map": row.get("map"),
            }
    return out


def step1_build_video_splits(out_dir: Path) -> Dict[str, Dict]:
    """Build per-dataset and merged video_split.json files."""
    logger.info("=== Step 1: building video splits ===")
    out_dir.mkdir(parents=True, exist_ok=True)

    collectors = {
        "nexar":         collect_nexar,
        "dad":           collect_dad,
        "dota":          collect_dota,
        "dada":          collect_dada,
        "adasto_critic": collect_adasto,
        "accident":      collect_accident,
    }

    merged = {}
    for name, fn in collectors.items():
        per_ds = fn()
        merged.update(per_ds)
        # per-dataset split file
        out_path = out_dir / f"{name}_split.json"
        out_path.write_text(json.dumps(per_ds, indent=2))
        logger.info(f"  {name}: {len(per_ds)} videos -> {out_path.name}")

    # merged
    merged_path = out_dir / "video_split.json"
    merged_path.write_text(json.dumps(merged, indent=2))
    logger.info(f"  merged: {len(merged)} videos -> {merged_path.name}")

    # summary stats
    print_split_summary(merged)
    write_summary_stats(merged, STATS_DIR)
    return merged


def print_split_summary(merged: Dict[str, Dict]) -> None:
    counts = defaultdict(lambda: defaultdict(lambda: defaultdict(int)))
    for v in merged.values():
        if v.get("excluded_from_main"):
            counts[v["source"]]["excluded_non_ego"][v["category"]] += 1
        else:
            counts[v["source"]][v["split"]][v["category"]] += 1

    lines = [
        "\n══════════ Split summary (video counts) ══════════",
        f"{'Source':<15} {'Split':<22} {'Category':<14} {'#Videos':>8}",
    ]
    grand_total = defaultdict(int)
    for src in sorted(counts.keys()):
        for split_name in sorted(counts[src].keys()):
            for cat in sorted(counts[src][split_name].keys()):
                n = counts[src][split_name][cat]
                lines.append(f"{src:<15} {split_name:<22} {cat:<14} {n:>8}")
                grand_total[split_name] += n
    lines.append("───────── totals per split ─────────")
    for sp in sorted(grand_total):
        lines.append(f"{'TOTAL':<15} {sp:<22} {'':<14} {grand_total[sp]:>8}")
    print("\n".join(lines))


def write_summary_stats(merged: Dict[str, Dict], stats_dir: Path) -> None:
    """Write per_source_video_count.csv with the same info."""
    stats_dir.mkdir(parents=True, exist_ok=True)
    rows = []
    counts = defaultdict(lambda: defaultdict(lambda: defaultdict(int)))
    for v in merged.values():
        sub = "excluded_non_ego" if v.get("excluded_from_main") else v["split"]
        counts[v["source"]][sub][v["category"]] += 1
    for src in sorted(counts):
        for split_name in sorted(counts[src]):
            for cat in sorted(counts[src][split_name]):
                rows.append({
                    "source": src,
                    "split": split_name,
                    "category": cat,
                    "n_videos": counts[src][split_name][cat],
                })
    import csv
    csv_path = stats_dir / "per_source_video_count.csv"
    with csv_path.open("w") as f:
        w = csv.DictWriter(f, fieldnames=list(rows[0].keys()))
        w.writeheader()
        w.writerows(rows)
    logger.info(f"  stats -> {csv_path}")


# ───────────────────────── main ─────────────────────────


# ═════════════════════ Step 2: per-frame action labels ═════════════════════

LABELS_DIR = BENCH_DIR / "labels"
DATA_DIR   = BENCH_DIR / "data"

SOURCE_FPS = {
    "nexar":         30.0,
    "dota":          10.0,
    "dad":           25.0,
    "dada":          30.0,
    "adasto_critic": 20.0,
    "accident":      20.0,
}
SILENT, OBSERVE, ALERT = 0, 1, 2
ACTION_NAME = {0: "SILENT", 1: "OBSERVE", 2: "ALERT"}

# Category remap for public-facing HF schema: drop ego/non-ego distinction.
def hf_category(raw_category: str) -> str:
    if raw_category in ("ego_positive", "non_ego"):
        return "positive"
    if raw_category == "safe_neg":
        return "negative"
    return "mixed"  # adasto_critic


def _probe_num_frames(video_path: Path) -> int:
    """Return num_frames using cv2 for .mp4, or listdir for frames-folder."""
    if video_path.is_dir():
        return len([f for f in video_path.iterdir()
                    if f.suffix.lower() in (".jpg", ".jpeg", ".png")])
    if video_path.suffix.lower() == ".mp4":
        import cv2
        cap = cv2.VideoCapture(str(video_path))
        n = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        cap.release()
        return n
    return 0


def _load_nexar_metadata() -> Dict[str, float]:
    """video_id -> time_of_event (seconds). Returns nan if missing/negative."""
    out: Dict[str, float] = {}
    import csv
    for folder in ("train/positive", "train/negative",
                   "test-public/positive", "test-public/negative",
                   "test-private/positive", "test-private/negative"):
        meta_csv = NEXAR_DIR / folder / "metadata.csv"
        if not meta_csv.exists():
            continue
        with meta_csv.open() as f:
            reader = csv.DictReader(f)
            for row in reader:
                fname = row.get("file_name", "")
                stem = Path(fname).stem
                if not stem:
                    continue
                t_event = row.get("time_of_event") or ""
                try:
                    out[f"nexar_{stem}"] = float(t_event) if t_event else float("nan")
                except ValueError:
                    out[f"nexar_{stem}"] = float("nan")
    return out


def _load_accident_metadata() -> Dict[str, dict]:
    """Kaggle ACCIDENT clip_name -> {t_takeover, duration, no_frames}"""
    import csv
    out: Dict[str, dict] = {}
    for csv_name in ("takeover_manifest_b50.csv", "takeover_manifest.csv"):
        p = CARLA_DIR / csv_name
        if not p.exists():
            continue
        with p.open() as f:
            for row in csv.DictReader(f):
                clip = row.get("clip")
                if clip and clip not in out:
                    out[clip] = {
                        "t_takeover": float(row.get("t_takeover", 0)),
                        "duration":   float(row.get("duration", 0)),
                        "no_frames":  int(row.get("no_frames", 0)),
                    }
    return out


def _load_dada_metadata() -> Dict[str, dict]:
    """folder_name -> {accident_time (frames), risky_time (frames)} from per-clip annotation.json."""
    out: Dict[str, dict] = {}
    for cat_dir in ("positive", "negative", "non-ego"):
        d = DADA_DIR / cat_dir
        if not d.exists():
            continue
        for sub in d.iterdir():
            if not sub.is_dir():
                continue
            ann = sub / "annotation.json"
            if not ann.exists():
                continue
            try:
                a = json.loads(ann.read_text())
                out[sub.name] = {
                    "accident_time": int(a.get("accident_time", -1)),
                    "risky_time":    int(a.get("risky_time", -1)),
                }
            except Exception:
                pass
    return out


def _build_labels_from_t_event(num_frames: int, fps: float,
                                t_event_s: float,
                                t_observe_window_s: float = 4.0,
                                t_alert_window_s: float = 2.0) -> List[int]:
    """Per-frame labels (0/1/2) given an event time in seconds.

    Convention: t_observe_window_s = 4.0 means OBSERVE starts 4s before event;
    t_alert_window_s = 2.0 means ALERT starts 2s before event.
    Post-event frames are SILENT (driver no longer needs alerting).
    """
    if t_event_s is None or not (t_event_s == t_event_s) or t_event_s < 0:
        return [SILENT] * num_frames
    t_alert_start = t_event_s - t_alert_window_s
    t_obs_start = t_event_s - t_observe_window_s
    labels = []
    for f in range(num_frames):
        t = f / fps
        if t >= t_event_s:
            labels.append(SILENT)
        elif t >= t_alert_start:
            labels.append(ALERT)
        elif t >= t_obs_start:
            labels.append(OBSERVE)
        else:
            labels.append(SILENT)
    return labels


def _labels_for_video(info: dict,
                      nexar_meta: Dict[str, float],
                      accident_meta: Dict[str, dict],
                      dada_meta: Dict[str, dict]) -> Optional[dict]:
    """Compute (num_frames, fps, labels, t_event_s) for one video."""
    src = info["source"]
    cat = info["category"]
    fps = SOURCE_FPS[src]
    video_path = ROOT / info["video_path"]
    is_positive = cat in ("ego_positive", "non_ego")  # both → "positive" for alerting

    try:
        if src == "nexar":
            num_frames = _probe_num_frames(video_path)
            if num_frames == 0:
                return None
            t_event = nexar_meta.get(info["video_id"], float("nan"))
            if cat == "safe_neg":
                t_event = float("nan")
            # BUG FIX: Nexar test-public / test-private positive videos are
            # CROPPED to ~10s ending just before the accident. The metadata
            # `time_of_event` refers to the ORIGINAL un-cropped video and is
            # therefore beyond our clip duration. For cropped test videos,
            # the event is effectively at the END of the clip (per Nexar
            # competition convention). Detect this case (clip duration <
            # metadata t_event) and override t_event to clip-end.
            if t_event == t_event and t_event > 0:
                clip_duration = num_frames / fps
                if t_event > clip_duration:
                    # Cropped video: event is at clip end (Nexar convention
                    # places accident in the final ~0.5s of test clips).
                    t_event = clip_duration  # end of clip
            labels = _build_labels_from_t_event(num_frames, fps, t_event)

        elif src == "dota":
            num_frames = info.get("num_frames") or _probe_num_frames(video_path / "images")
            anomaly_start = info.get("anomaly_start")  # in frames
            t_event = anomaly_start / fps if anomaly_start else float("nan")
            labels = _build_labels_from_t_event(num_frames, fps, t_event)

        elif src == "dad":
            # All DAD videos are 4s @ 25fps; accident at the END (t=4.0)
            num_frames = 100
            t_event = 4.0 if is_positive else float("nan")
            labels = _build_labels_from_t_event(num_frames, fps, t_event)

        elif src == "dada":
            num_frames = _probe_num_frames(video_path)
            if num_frames == 0:
                return None
            meta = dada_meta.get(video_path.name, {})
            acc_f = meta.get("accident_time", -1)
            t_event = acc_f / fps if acc_f and acc_f > 0 else float("nan")
            if cat == "safe_neg":
                t_event = float("nan")
            labels = _build_labels_from_t_event(num_frames, fps, t_event)

        elif src == "adasto_critic":
            # ADAS-TO-Critic clips are uniformly 20s @ 20fps = 400 frames; t_takeover=10s
            num_frames = 400
            t_event = info.get("t_takeover_s", 10.0)
            labels = _build_labels_from_t_event(num_frames, fps, t_event)

        elif src == "accident":
            cm = accident_meta.get(Path(info["video_path"]).stem, {})
            num_frames = cm.get("no_frames") or _probe_num_frames(video_path)
            if num_frames == 0:
                return None
            t_event = cm.get("t_takeover", info.get("t_takeover_s", float("nan")))
            labels = _build_labels_from_t_event(num_frames, fps, t_event)
        else:
            return None
    except Exception as e:
        logger.warning(f"label compute failed for {info['video_id']}: {e}")
        return None

    return {
        "num_frames": num_frames,
        "fps":        fps,
        "t_event_s":  None if not (t_event == t_event) else float(t_event),
        "labels":     labels,
    }


def step2_per_frame_labels(out_dir: Path) -> None:
    """Generate per-frame action labels per video for all 4 splits (train/val/test/extra)."""
    logger.info("=== Step 2: per-frame action labels ===")
    out_dir.mkdir(parents=True, exist_ok=True)
    video_split = json.loads((MANIFEST_DIR / "video_split.json").read_text())

    logger.info("  loading per-source metadata caches...")
    nexar_meta    = _load_nexar_metadata()
    accident_meta = _load_accident_metadata()
    dada_meta     = _load_dada_metadata()
    logger.info(f"    nexar:    {len(nexar_meta)} entries")
    logger.info(f"    accident: {len(accident_meta)} entries")
    logger.info(f"    dada:     {len(dada_meta)} entries")

    per_split = defaultdict(list)
    fail_count = defaultdict(int)
    total = len(video_split)
    for i, (vid_id, info) in enumerate(video_split.items()):
        if i % 500 == 0:
            logger.info(f"  [{i}/{total}] processing...")
        split = info["split"]
        if split == "excluded_non_ego":
            continue
        result = _labels_for_video(info, nexar_meta, accident_meta, dada_meta)
        if result is None:
            fail_count[info["source"]] += 1
            continue
        record = {
            "video_id":     vid_id,
            "source":       info["source"],
            "split":        split,
            "category":     hf_category(info["category"]),     # public-facing
            "raw_category": info["category"],                  # internal
            "video_path":   info["video_path"],
            "native_split": info.get("native_split"),
            **result,
        }
        # add source-specific extras
        for k in ("anomaly_class", "anomaly_start", "anomaly_end",
                  "t_takeover_s", "accident_type"):
            if k in info:
                record[k] = info[k]
        per_split[split].append(record)

    for split, records in per_split.items():
        out_path = out_dir / f"{split}_perframe.json"
        out_path.write_text(json.dumps(
            {"split": split, "n_videos": len(records), "samples": records}))
        # action distribution sanity
        cnt = Counter(a for r in records for a in r["labels"])
        n_total = sum(cnt.values()) or 1
        dist = {ACTION_NAME[k]: f"{cnt[k]/n_total:.3f}" for k in (SILENT, OBSERVE, ALERT)}
        logger.info(f"  {split}: {len(records)} videos -> {out_path.name}  action_dist={dist}")
    if fail_count:
        logger.warning(f"  failed videos (skipped): {dict(fail_count)}")


# ═════════════════════ Step 3: tick-level parquet ═════════════════════

def step3_tick_parquet(out_dir: Path,
                        win_frames: int = 8,
                        tick_hz: float = 1.0) -> None:
    """Sliding 8-frame window at 1Hz tick rate -> Parquet per split."""
    logger.info("=== Step 3: tick-level parquet ===")
    out_dir.mkdir(parents=True, exist_ok=True)
    try:
        import pyarrow as pa
        import pyarrow.parquet as pq
    except ImportError:
        logger.error("pyarrow not installed. pip install pyarrow")
        return

    for label_path in sorted(LABELS_DIR.glob("*_perframe.json")):
        split = label_path.stem.replace("_perframe", "")
        doc = json.loads(label_path.read_text())
        ticks = []
        for vid in doc["samples"]:
            n = vid["num_frames"]
            fps = vid["fps"]
            stride = int(round(fps / tick_hz))  # 1 tick per second
            t_event = vid.get("t_event_s")
            for end_f in range(win_frames, n + 1, stride):
                frame_idx = list(range(end_f - win_frames, end_f))
                # Tick label = label at last frame in window
                last_f = end_f - 1
                tick_lbl = vid["labels"][last_f]
                # tta_raw: positive = (event_frame - last_f) / fps; nan if no event
                if t_event is None:
                    tta_raw = -1.0
                else:
                    tta_raw = float(t_event - last_f / fps)
                ticks.append({
                    "video_id":    vid["video_id"],
                    "source":      vid["source"],
                    "category":    vid["category"],
                    "split":       split,
                    "frame_indices": frame_idx,
                    "n_frames":    n,
                    "fps":         fps,
                    "tta_raw":     tta_raw,
                    "tick_label":  tick_lbl,
                    "video_path":  vid["video_path"],
                })
        if not ticks:
            logger.warning(f"  {split}: 0 ticks generated (empty?)")
            continue
        # Write parquet
        out_path = out_dir / f"{split}.parquet"
        table = pa.Table.from_pylist(ticks)
        pq.write_table(table, out_path, compression="snappy")
        cnt = Counter(t["tick_label"] for t in ticks)
        n_t = len(ticks)
        dist = {ACTION_NAME[k]: f"{cnt[k]/n_t:.3f}" for k in (SILENT, OBSERVE, ALERT)}
        logger.info(f"  {split}: {n_t} ticks -> {out_path.name}  tick_dist={dist}")


# ═════════════════════ Step 4: HF loader + dataset card ═════════════════════

LOADER_PY_TEMPLATE = '''"""VLAlert-Bench: unified driving-alert benchmark.

This loader exposes per-tick records (1Hz sliding window over 8 frames) with
SILENT/OBSERVE/ALERT action targets. Videos are NOT redistributed — users must
download source datasets from their original providers (see README) and pass
local paths to from_local_video() to materialize frames.

Splits:
  - train, val, test:           in-domain (Nexar + DoTA + DAD + DADA-2000)
  - extra_val_adasto:           held-out OOD (ADAS-TO-Critic, full corpus)
  - extra_val_accident:         held-out OOD (Kaggle ACCIDENT @ CVPR 2026)
"""
import datasets
import json
import os

_CITATION = """@article{wang2026vlalert,
  title={VLAlert-X: A Vision-Language POMDP for Driving-Alert Decisions},
  author={Wang, Anonymous and others},
  year={2026}
}"""

_DESCRIPTION = """VLAlert-Bench unifies 6 driving-event datasets (Nexar Collision,
DoTA, DAD, DADA-2000, ADAS-TO-Critic, Kaggle ACCIDENT @ CVPR 2026) into
per-tick records with 3-way action labels (SILENT/OBSERVE/ALERT). Five
splits: train / val / test / extra_val_adasto / extra_val_accident.
Annotations are released here; source videos remain under their original
licenses (ADAS-TO-Critic mp4s are co-hosted in this repo)."""

_HOMEPAGE = "https://huggingface.co/datasets/AsianPlayer/VLAlert"
_LICENSE  = "Annotations: CC-BY-4.0. Source videos: see README per-source licenses."


class VLAlertBenchConfig(datasets.BuilderConfig):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)


class VLAlertBench(datasets.GeneratorBasedBuilder):
    VERSION = datasets.Version("1.0.0")
    BUILDER_CONFIGS = [VLAlertBenchConfig(name="default", version=VERSION,
                                          description="Default per-tick view.")]

    def _info(self):
        return datasets.DatasetInfo(
            description=_DESCRIPTION,
            features=datasets.Features({
                "video_id":      datasets.Value("string"),
                "source":        datasets.ClassLabel(names=["nexar","dota","dad","dada","adasto_critic","accident"]),
                "category":      datasets.ClassLabel(names=["positive","negative","mixed"]),
                "split":         datasets.Value("string"),
                "frame_indices": datasets.Sequence(datasets.Value("int32")),
                "n_frames":      datasets.Value("int32"),
                "fps":           datasets.Value("float32"),
                "tta_raw":       datasets.Value("float32"),
                "tick_label":    datasets.ClassLabel(names=["SILENT","OBSERVE","ALERT"]),
                "video_path":    datasets.Value("string"),
            }),
            supervised_keys=None,
            homepage=_HOMEPAGE,
            license=_LICENSE,
            citation=_CITATION,
        )

    def _split_generators(self, dl_manager):
        data_dir = os.path.join(self.config.data_dir or "data")
        return [
            datasets.SplitGenerator(name=datasets.Split.TRAIN,
                                     gen_kwargs={"path": os.path.join(data_dir, "train.parquet")}),
            datasets.SplitGenerator(name=datasets.Split.VALIDATION,
                                     gen_kwargs={"path": os.path.join(data_dir, "val.parquet")}),
            datasets.SplitGenerator(name=datasets.Split.TEST,
                                     gen_kwargs={"path": os.path.join(data_dir, "test.parquet")}),
            datasets.SplitGenerator(name="extra_val_adasto",
                                     gen_kwargs={"path": os.path.join(data_dir, "extra_val_adasto.parquet")}),
            datasets.SplitGenerator(name="extra_val_accident",
                                     gen_kwargs={"path": os.path.join(data_dir, "extra_val_accident.parquet")}),
        ]

    def _generate_examples(self, path):
        import pyarrow.parquet as pq
        table = pq.read_table(path)
        for i, row in enumerate(table.to_pylist()):
            yield i, row
'''


def step4_hf_loader(out_dir: Path) -> None:
    """Write vlalert_bench.py loader + dataset_infos.json metadata."""
    logger.info("=== Step 4: HF loader + dataset card ===")
    (out_dir / "vlalert_bench.py").write_text(LOADER_PY_TEMPLATE)
    logger.info(f"  loader -> vlalert_bench.py")
    # dataset_infos.json (lightweight; real one auto-generated by hf datasets)
    info = {
        "default": {
            "description": "VLAlert-Bench unified driving-alert benchmark.",
            "citation": "Wang et al. 2026",
            "homepage": "https://huggingface.co/datasets/AsianPlayer/VLAlert",
            "license":  "Annotations CC-BY-4.0; sources per README.",
            "features": {
                "video_id":      "string",
                "source":        "ClassLabel(nexar,dota,dad,dada,adasto_critic,accident)",
                "category":      "ClassLabel(positive,negative,mixed)",
                "frame_indices": "Sequence(int32,8)",
                "tta_raw":       "float32",
                "tick_label":    "ClassLabel(SILENT,OBSERVE,ALERT)",
            },
        }
    }
    (out_dir / "dataset_infos.json").write_text(json.dumps(info, indent=2))
    logger.info(f"  dataset_infos.json")


# ═════════════════════ Step 5: leakage verify + smoke test ═════════════════════

def step5_verify(out_dir: Path) -> None:
    """Cross-split video_id leakage check + parquet smoke load."""
    logger.info("=== Step 5: leakage verify + smoke test ===")
    out_dir.mkdir(parents=True, exist_ok=True)
    video_split = json.loads((MANIFEST_DIR / "video_split.json").read_text())
    splits = defaultdict(set)
    for vid_id, info in video_split.items():
        splits[info["split"]].add(vid_id)
    # Pairwise leakage across all 5 in-corpus splits
    in_corpus = ["train", "val", "test", "extra_val_adasto", "extra_val_accident"]
    pairs = [(a, b) for i, a in enumerate(in_corpus)
                    for b in in_corpus[i + 1:]]
    leakage = {}
    for a, b in pairs:
        overlap = splits[a] & splits[b]
        leakage[f"{a}__{b}"] = {"n_overlap": len(overlap),
                                  "examples": list(overlap)[:5]}
    # Smoke: try loading each parquet, sample first 3 rows
    smoke = {}
    try:
        import pyarrow.parquet as pq
        for parquet_path in sorted(DATA_DIR.glob("*.parquet")):
            t = pq.read_table(parquet_path)
            smoke[parquet_path.stem] = {
                "n_rows":  t.num_rows,
                "columns": t.column_names,
                "first_video_ids": t.column("video_id").to_pylist()[:3],
            }
    except Exception as e:
        smoke["error"] = str(e)
    report = {"leakage": leakage, "smoke_load": smoke,
              "max_leakage": max((v["n_overlap"] for v in leakage.values()), default=0)}
    out_path = out_dir / "leakage_report.json"
    out_path.write_text(json.dumps(report, indent=2))
    logger.info(f"  report -> {out_path}")
    if report["max_leakage"] == 0:
        logger.info("  ✅ Zero video-id leakage across splits")
    else:
        logger.warning(f"  ⚠️ Leakage detected (max {report['max_leakage']}); see report.")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--step", choices=["1", "2", "3", "4", "5", "all"],
                    default="1")
    ap.add_argument("--out", type=Path, default=BENCH_DIR)
    args = ap.parse_args()

    if args.step in ("1", "all"):
        step1_build_video_splits(args.out / "manifest")
    if args.step in ("2", "all"):
        step2_per_frame_labels(args.out / "labels")
    if args.step in ("3", "all"):
        step3_tick_parquet(args.out / "data")
    if args.step in ("4", "all"):
        step4_hf_loader(args.out)
    if args.step in ("5", "all"):
        step5_verify(args.out / "stats")


if __name__ == "__main__":
    main()