#!/usr/bin/env python3
"""Build per-frame CoT+Belief JSONL for DoTA clips.

DoTA structure:
  DoTA/frames/{clip}/images/000000.jpg, 000001.jpg, ...
  DoTA/annotations/{clip}.json         — has anomaly_start/end (frame idx), ego_involve,
                                         accident_name, night, per-frame object list
  DoTA/train_split.txt, val_split.txt

We emit one JSONL record per clip with:
  {
    "id":         str (clip name),
    "label":      int (1 = ego-involved anomaly, 0 = normal segment BEFORE anomaly),
    "video_path": str (absolute path to frames dir),
    "cot":        {"scene": ..., "critical_objects": [...], "threat_analysis": ...},
    "belief":     {"action": str,
                   "actions_per_frame": [T str],
                   "frame_times_sec":   [T float],
                   "time_of_event_sec": float,
                   "alert_s":           float,
                   "observe_s":         float}
  }

CoT content is a rule-based placeholder (no teacher model). Replace in a later
pass if the budget allows running GPT-4o over DoTA val.

Usage:
  python -m training.VLA.build_dota_cot \
      --dota_root DoTA \
      --split val \
      --n_frames 8 \
      --fps 10 \
      --alert_s 0.5 --observe_s 2.5 \
      --output data/vla_cot/dota_val_perframe.jsonl
"""
from __future__ import annotations

import argparse
import json
from pathlib import Path
from typing import Any, Dict, List, Optional

import numpy as np

from training.VLA.augment_cot_with_belief import (
    derive_clip_action, per_frame_actions,
)


# ── tiny helpers ──────────────────────────────────────────────────────────
NIGHT_STR = {True: "night", False: "daytime"}


def _threat_from_accident(accident_name: str, ego_involve: bool) -> str:
    nm = (accident_name or "").replace("_", " ").strip()
    if not nm:
        return "Possible traffic anomaly ahead; trajectory appears unstable."
    if ego_involve:
        return (f"Ego vehicle is at elevated collision risk: the scene shows "
                f"a '{nm}' pattern developing in the ego lane.")
    return (f"Nearby agents exhibit a '{nm}' pattern; ego not directly involved "
            "but secondary risk cannot be ruled out.")


def _scene_str(accident_name: str, night: bool) -> str:
    tag = "night" if night else "daytime"
    nm = (accident_name or "unknown").replace("_", " ")
    return f"Dashcam view, {tag}; anomaly pattern: {nm}."


def _critical_from_labels(labels: List[Dict[str, Any]], anomaly_start: int,
                          n_keep: int = 4) -> List[str]:
    """Pick up to n_keep distinct object classes that appear near anomaly_start."""
    names: Dict[str, int] = {}
    if not labels:
        return []
    lo = max(0, anomaly_start - 5)
    hi = min(len(labels), anomaly_start + 3)
    for lab in labels[lo:hi]:
        for obj in lab.get("objects", []) or []:
            cat = obj.get("category") or obj.get("label") or obj.get("accident_name")
            if cat and cat != "normal":
                names[cat] = names.get(cat, 0) + 1
    if not names:
        return []
    ranked = sorted(names.items(), key=lambda kv: -kv[1])
    return [k for k, _ in ranked[:n_keep]]


def _uniform_indices(total: int, n: int) -> List[int]:
    return np.linspace(0, total - 1, n).round().astype(int).tolist()


def build_record(
    clip_name: str,
    ann: Dict[str, Any],
    frames_root: Path,
    n_frames: int,
    fps: float,
    alert_s: float,
    observe_s: float,
    event_anchored: bool = True,
    lookback_s: float = 3.0,
    post_margin_s: float = 0.0,
    teacher_cot: Optional[Dict[str, Any]] = None,
) -> Optional[Dict[str, Any]]:
    total = int(ann.get("num_frames", 0))
    if total <= 0:
        return None
    ego = bool(ann.get("ego_involve", False))
    night = bool(ann.get("night", False))
    acc_name = str(ann.get("accident_name") or ann.get("anomaly_class") or "")
    anomaly_start = int(ann.get("anomaly_start", -1))

    # label = 1 iff ego-involved anomaly clip (a positive sample for POMDP SFT)
    # non-ego anomaly clips are excluded (they add noise; ego_involve handles them).
    if not ego:
        return None
    label = 1

    # time of event = anomaly_start / fps (seconds from clip start)
    toe = float(anomaly_start) / float(fps) if anomaly_start >= 0 else None

    if event_anchored and toe is not None:
        end_s   = min(float(total - 1) / fps, toe + post_margin_s)
        start_s = max(0.0, end_s - (lookback_s + post_margin_s))
        times = np.linspace(start_s, end_s, n_frames)
        idx = [int(round(t * fps)) for t in times]
        sampling = "event_anchored"
    else:
        idx = _uniform_indices(total, n_frames)
        sampling = "uniform"
    idx = [max(0, min(total - 1, int(i))) for i in idx]
    frame_times = [float(i) / float(fps) for i in idx]
    actions_pf = per_frame_actions(label=label,
                                   time_of_event=toe,
                                   frame_times=frame_times,
                                   alert_s=alert_s,
                                   observe_s=observe_s)

    # clip-level action: if any ALERT → ALERT; elif any OBSERVE → OBSERVE; else SILENT
    if "ALERT" in actions_pf:
        clip_action = "ALERT"
    elif "OBSERVE" in actions_pf:
        clip_action = "OBSERVE"
    else:
        clip_action = "SILENT"

    if teacher_cot is not None and isinstance(teacher_cot, dict) \
            and teacher_cot.get("scene") and teacher_cot.get("threat_analysis"):
        cot = {
            "scene":            str(teacher_cot.get("scene", "")).strip(),
            "critical_objects": list(teacher_cot.get("critical_objects", []) or []),
            "threat_analysis":  str(teacher_cot.get("threat_analysis", "")).strip(),
            "source":           "gpt_teacher",
        }
    else:
        cot = {
            "scene":            _scene_str(acc_name, night),
            "critical_objects": _critical_from_labels(ann.get("labels", []) or [],
                                                     anomaly_start),
            "threat_analysis":  _threat_from_accident(acc_name, ego),
            "source":           "rule_template",
        }
    belief = {
        "action":            clip_action,
        "tta_sec":           round(toe, 3) if toe is not None else -1.0,
        "actions_per_frame": actions_pf,
        "frame_indices":     idx,
        "frame_times_sec":   [round(t, 3) for t in frame_times],
        "time_of_event_sec": round(toe, 3) if toe is not None else -1.0,
        "alert_s":           alert_s,
        "observe_s":         observe_s,
        "sampling":          sampling,
        "total_frames":      total,
        "fps":               round(float(fps), 3),
        "lookback_s":        lookback_s,
        "post_margin_s":     post_margin_s,
    }

    return {
        "id":         f"dota_{clip_name}",
        "label":      label,
        "video_path": str(frames_root / clip_name / "images"),
        "cot":        cot,
        "belief":     belief,
        "source":     "dota",
        "accident_name": acc_name,
    }


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--dota_root", default="DoTA")
    ap.add_argument("--split", choices=["train", "val"], default="val")
    ap.add_argument("--n_frames",  type=int,   default=8)
    ap.add_argument("--fps",       type=float, default=10.0)
    ap.add_argument("--alert_s",   type=float, default=0.5)
    ap.add_argument("--observe_s", type=float, default=2.5)
    ap.add_argument("--output",    required=True)
    ap.add_argument("--event_anchored", action="store_true", default=True,
                    help="Sample T frames in [event - lookback, event + post_margin]")
    ap.add_argument("--no_event_anchored", dest="event_anchored",
                    action="store_false")
    ap.add_argument("--lookback_s",    type=float, default=3.0)
    ap.add_argument("--post_margin_s", type=float, default=0.0)
    ap.add_argument("--teacher_json",  type=str,   default="",
                    help="Optional JSON mapping clip_name -> {cot, usage, ...} "
                         "from GPT distillation; overrides rule template when present")
    ap.add_argument("--limit",     type=int,   default=0,
                    help="If >0, only emit first N records (smoke test)")
    args = ap.parse_args()

    root = Path(args.dota_root).resolve()
    split_file = root / f"{args.split}_split.txt"
    ann_dir    = root / "annotations"
    frames_root = root / "frames"

    with open(split_file) as f:
        clips = [ln.strip() for ln in f if ln.strip()]
    print(f"[dota] {args.split}: {len(clips)} clip ids")

    teacher_map: Dict[str, Dict[str, Any]] = {}
    if args.teacher_json:
        tp = Path(args.teacher_json)
        if tp.exists():
            raw = json.loads(tp.read_text())
            for k, v in raw.items():
                if isinstance(v, dict) and v.get("cot"):
                    teacher_map[k] = v["cot"]
            print(f"[dota] teacher CoT loaded: {len(teacher_map)} clips from {tp}")
        else:
            print(f"[dota] WARN: teacher_json not found: {tp}")

    out_path = Path(args.output)
    out_path.parent.mkdir(parents=True, exist_ok=True)

    n_emitted = 0
    n_missing = 0
    n_no_ego = 0
    with open(out_path, "w") as f_out:
        for clip_name in clips:
            ann_path = ann_dir / f"{clip_name}.json"
            if not ann_path.exists():
                n_missing += 1
                continue
            with open(ann_path) as f:
                ann = json.load(f)
            frames_dir = frames_root / clip_name / "images"
            if not frames_dir.exists():
                n_missing += 1
                continue
            rec = build_record(
                clip_name=clip_name,
                ann=ann,
                frames_root=frames_root,
                n_frames=args.n_frames,
                fps=args.fps,
                alert_s=args.alert_s,
                observe_s=args.observe_s,
                event_anchored=args.event_anchored,
                lookback_s=args.lookback_s,
                post_margin_s=args.post_margin_s,
                teacher_cot=teacher_map.get(clip_name),
            )
            if rec is None:
                n_no_ego += 1
                continue
            f_out.write(json.dumps(rec) + "\n")
            n_emitted += 1
            if args.limit and n_emitted >= args.limit:
                break

    print(f"[dota] emitted {n_emitted} ego-involved clips; "
          f"skipped non-ego={n_no_ego}  missing={n_missing}")
    print(f"[dota] saved -> {out_path}")


if __name__ == "__main__":
    main()