#!/usr/bin/env python3
"""
Generate DPO preference-pair manifests from SFT video manifests.

Pair logic
----------
For each ego_positive video:
  chosen   : windows where TTA ∈ [CHOSEN_TTA_MIN, CHOSEN_TTA_MAX]  → model SHOULD alert
  rejected : windows where TTA  > REJECTED_EARLY_MIN               → too early to alert
             windows where TTA  < REJECTED_LATE_MAX                → too late (useless)

For each safe_neg / non_ego video:
  These are NEVER-alert windows.  They are paired cross-video against
  a randomly sampled ego_pos chosen window (same source preferred).

Output
------
data/dpo_pairs/
  nexar_train.json
  dada_train.json
  nexar_val.json
  dada_val.json

Usage
-----
cd PROJECT_ROOT
python -m training.DPO.make_dpo_pairs \
    --manifest_dir data/sft_manifests \
    --out_dir      data/dpo_pairs
"""

from __future__ import annotations

import argparse
import json
import logging
import random
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("DPO.make_pairs")

# ── constants (must match SFT dataset.py) ────────────────────────────────────
FRAME_RATE       = 20
WINDOW_LEN       = 40    # 2.0 s
SAMPLE_RATE      = 4     # keep every 4th frame inside window
MAX_FRAMES       = 8

# Alert timing targets
CHOSEN_TTA_MIN   = 1.5   # seconds  (sweet-spot alert window)
CHOSEN_TTA_MAX   = 5.0
CHOSEN_TTA_STEPS = [2.0, 2.5, 3.0, 3.5, 4.0, 4.5]   # chosen TTA values to sample

REJECTED_EARLY_MIN   = 5.5                             # too early
REJECTED_EARLY_STEPS = [6.0, 7.0, 8.0, 9.0]
REJECTED_LATE_MAX    = 1.0                             # too late (reaction impossible)
REJECTED_LATE_STEPS  = [0.5, 1.0]

RANDOM_SEED = 42


# ─────────────────────────────────────────────────────────────────────────────
# Frame index helpers
# ─────────────────────────────────────────────────────────────────────────────

def build_window(
    window_end: int,
    num_frames: int,
    window_len: int = WINDOW_LEN,
    sample_rate: int = SAMPLE_RATE,
    max_frames: int = MAX_FRAMES,
) -> Optional[List[int]]:
    """Return sampled frame indices for a window ending at `window_end` (exclusive).
    Returns None if the window falls outside [0, num_frames)."""
    window_start = window_end - window_len
    if window_start < 0 or window_end > num_frames:
        return None
    indices = list(range(window_start, window_end, sample_rate))[:max_frames]
    if not indices:
        return None
    return indices


def window_entry(
    source_dir: str,
    frame_indices: List[int],
    window_end: int,
    tta_true: float,
    label: str,
    metadata: dict,
) -> dict:
    return {
        "source_dir":    source_dir,
        "frame_indices": frame_indices,
        "window_end":    window_end,
        "tta_true":      round(tta_true, 3),
        "label":         label,
        "metadata":      metadata,
    }


# ─────────────────────────────────────────────────────────────────────────────
# Pair generators
# ─────────────────────────────────────────────────────────────────────────────

def generate_ego_pos_pairs(video: dict) -> List[dict]:
    """Return (chosen, rejected) pairs for a single ego_pos video."""
    src       = video["source_dir"]
    nf        = video["num_frames"]
    af        = video["accident_frame"]
    meta      = video["metadata"]
    vid       = video["video_id"]
    source    = video["source"]

    if af is None:
        return []

    chosen_windows:   List[Tuple[float, List[int], int]] = []  # (tta, indices, w_end)
    rejected_windows: List[Tuple[float, str, List[int], int]] = []  # (tta, label, ...)

    # ── chosen windows ────────────────────────────────────────────────────────
    for tta in CHOSEN_TTA_STEPS:
        w_end = af - round(tta * FRAME_RATE)
        idxs  = build_window(w_end, nf)
        if idxs is not None:
            chosen_windows.append((tta, idxs, w_end))

    # ── rejected early ────────────────────────────────────────────────────────
    for tta in REJECTED_EARLY_STEPS:
        w_end = af - round(tta * FRAME_RATE)
        idxs  = build_window(w_end, nf)
        if idxs is not None:
            rejected_windows.append((tta, "too_early", idxs, w_end))

    # ── rejected late ─────────────────────────────────────────────────────────
    for tta in REJECTED_LATE_STEPS:
        w_end = af - round(tta * FRAME_RATE)
        idxs  = build_window(w_end, nf)
        if idxs is not None:
            rejected_windows.append((tta, "too_late", idxs, w_end))

    if not chosen_windows or not rejected_windows:
        return []

    pairs = []
    for c_tta, c_idxs, c_wend in chosen_windows:
        for r_tta, r_label, r_idxs, r_wend in rejected_windows:
            pairs.append({
                "pair_id":    f"{vid}_c{c_tta}_r{r_tta}_{r_label}",
                "video_id":   vid,
                "source":     source,
                "pair_type":  "timing",
                "chosen":     window_entry(src, c_idxs, c_wend, c_tta, "timely_alert", meta),
                "rejected":   window_entry(src, r_idxs, r_wend, r_tta, r_label, meta),
            })
    return pairs


def generate_neg_windows(video: dict) -> List[dict]:
    """Return 'never-alert' window entries for safe_neg / non_ego videos."""
    src    = video["source_dir"]
    nf     = video["num_frames"]
    meta   = video["metadata"]
    cat    = video["category"]

    # Sample windows from the middle third of the video
    start = nf // 3
    end   = 2 * nf // 3
    entries = []
    stride = max(1, (end - start) // 3)
    for w_end in range(start + WINDOW_LEN, end, stride):
        idxs = build_window(w_end, nf)
        if idxs is not None:
            entries.append(window_entry(src, idxs, w_end, tta_true=999.0, label=cat, metadata=meta))
    return entries[:3]  # cap at 3 windows per video


# ─────────────────────────────────────────────────────────────────────────────
# Manifest processing
# ─────────────────────────────────────────────────────────────────────────────

def process_manifests(
    manifests:      List[Path],
    split:          str,
    rng:            random.Random,
    max_cross_pairs: int = 3,
) -> List[dict]:
    """Build all DPO pairs from a list of manifest files."""
    all_videos: List[dict] = []
    for m in manifests:
        if not m.exists():
            logger.warning(f"Manifest not found: {m}")
            continue
        with open(m) as f:
            data = json.load(f)
        vids = data.get("videos", [])
        logger.info(f"  {m.name}: {len(vids)} videos")
        all_videos.extend(vids)

    ego_pos  = [v for v in all_videos if v["category"] == "ego_positive"]
    neg_vids = [v for v in all_videos if v["category"] in ("safe_neg", "non_ego")]

    pairs: List[dict] = []

    # ── within-video timing pairs (ego_pos) ───────────────────────────────────
    for v in ego_pos:
        pairs.extend(generate_ego_pos_pairs(v))

    # ── cross-type pairs (neg window vs chosen ego_pos window) ───────────────
    if ego_pos and neg_vids:
        # Build pool of chosen windows from ego_pos (for cross-pairing)
        chosen_pool: Dict[str, List[dict]] = {}   # source → [chosen_entry]
        for v in ego_pos:
            sub_pairs = generate_ego_pos_pairs(v)
            for p in sub_pairs:
                src = v["source"]
                chosen_pool.setdefault(src, []).append(
                    (v["video_id"], p["chosen"])
                )

        for nv in neg_vids:
            neg_entries = generate_neg_windows(nv)
            if not neg_entries:
                continue
            src = nv["source"]
            pool = chosen_pool.get(src, [])
            if not pool:
                pool = [item for items in chosen_pool.values() for item in items]
            if not pool:
                continue
            for ne in neg_entries[:max_cross_pairs]:
                vid_c, c_entry = rng.choice(pool)
                pairs.append({
                    "pair_id":    f"cross_{nv['video_id']}_{vid_c}",
                    "video_id":   nv["video_id"],
                    "source":     nv["source"],
                    "pair_type":  "category",
                    "chosen":     c_entry,
                    "rejected":   ne,
                })

    logger.info(f"  Split={split}: {len(pairs)} total pairs "
                f"({sum(1 for p in pairs if p['pair_type']=='timing')} timing, "
                f"{sum(1 for p in pairs if p['pair_type']=='category')} category)")
    return pairs


# ─────────────────────────────────────────────────────────────────────────────
# Main
# ─────────────────────────────────────────────────────────────────────────────

def main():
    parser = argparse.ArgumentParser("make_dpo_pairs")
    parser.add_argument("--manifest_dir", default="data/sft_manifests")
    parser.add_argument("--out_dir",      default="data/dpo_pairs")
    parser.add_argument("--seed",         type=int, default=RANDOM_SEED)
    args = parser.parse_args()

    mdir = Path(args.manifest_dir)
    odir = Path(args.out_dir)
    odir.mkdir(parents=True, exist_ok=True)
    rng  = random.Random(args.seed)

    splits = {
        "nexar_train": [mdir / "nexar_train.json"],
        "dada_train":  [mdir / "dada_pos_train.json",
                        mdir / "dada_noneego_train.json",
                        mdir / "dada_neg_train.json"],
        "nexar_val":   [mdir / "nexar_val.json"],
        "dada_val":    [mdir / "dada_pos_val.json",
                        mdir / "dada_noneego_val.json"],
    }

    for name, manifests in splits.items():
        split = "train" if "train" in name else "val"
        logger.info(f"\nProcessing {name} ...")
        pairs = process_manifests(manifests, split, rng)
        if split == "train":
            rng.shuffle(pairs)
        out_path = odir / f"{name}.json"
        with open(out_path, "w") as f:
            json.dump({"name": name, "split": split,
                       "num_pairs": len(pairs), "pairs": pairs}, f)
        logger.info(f"  Saved {len(pairs)} pairs → {out_path}")

    logger.info("\n✅ DPO pair manifests generated.")


if __name__ == "__main__":
    main()