#!/usr/bin/env python3
"""
Extract frames from DAD (Dashcam Accident Dataset) MP4 files.
=============================================================
DAD has no timing annotations — just positive/negative video splits.
  positive/: 455 crash videos  (no exact accident timestamp)
  negative/: 829 safe driving  (no annotations needed)

Frames are extracted at 20 fps and saved as:
  data/pretrain_v2/dad_frames/{positive,negative}/{vid_id}/000.jpg ...

Each video also gets a minimal annotation.json:
  { "split": "positive"|"negative",
    "n_frames": int,
    "accident": bool,
    "fps": 20 }

Usage:
    python training/pretrain_v2/prepare_dad_frames.py [--workers 4]
"""

import argparse
import json
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path

import cv2

DAD_VIDEOS_DIR = "PROJECT_ROOT/DAD/videos/training"
OUTPUT_BASE    = "PROJECT_ROOT/data/pretrain_v2/dad_frames"
DAD_TEST_VIDEOS_DIR = "PROJECT_ROOT/DAD/videos/testing"
OUTPUT_BASE_TEST    = "PROJECT_ROOT/data/pretrain_v2/dad_frames_test"
TARGET_FPS     = 20.0
MAX_FRAMES     = 400   # cap at 400 frames (~20s at 20fps)


def extract_video(args):
    """Extract one video. Returns (video_id, n_frames, error|None)."""
    vid_path, out_dir, split = args
    vid_id = vid_path.stem
    out_dir.mkdir(parents=True, exist_ok=True)

    ann_path = out_dir / "annotation.json"
    if ann_path.exists():
        # Already extracted — skip
        existing = json.loads(ann_path.read_text())
        return vid_id, existing.get("n_frames", 0), None

    cap = cv2.VideoCapture(str(vid_path))
    if not cap.isOpened():
        return vid_id, 0, f"Cannot open {vid_path}"

    orig_fps = cap.get(cv2.CAP_PROP_FPS) or TARGET_FPS
    interval = max(1, int(round(orig_fps / TARGET_FPS)))

    raw_idx = 0
    saved   = 0
    while saved < MAX_FRAMES:
        ret, frame = cap.read()
        if not ret:
            break
        if raw_idx % interval == 0:
            cv2.imwrite(str(out_dir / f"{saved:03d}.jpg"), frame)
            saved += 1
        raw_idx += 1
    cap.release()

    if saved == 0:
        return vid_id, 0, f"No frames extracted from {vid_path}"

    ann_path.write_text(json.dumps({
        "split":    split,
        "accident": (split == "positive"),
        "n_frames": saved,
        "fps":      TARGET_FPS,
    }, indent=2), encoding="utf-8")

    return vid_id, saved, None


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--workers", type=int, default=4)
    parser.add_argument("--split", choices=["training", "testing"], default="training",
                        help="DAD split to extract: training or testing")
    args = parser.parse_args()

    if args.split == "testing":
        dad_root   = Path(DAD_TEST_VIDEOS_DIR)
        output_base = Path(OUTPUT_BASE_TEST)
    else:
        dad_root   = Path(DAD_VIDEOS_DIR)
        output_base = Path(OUTPUT_BASE)

    tasks = []
    for split in ("positive", "negative"):
        vid_dir = dad_root / split
        if not vid_dir.exists():
            print(f"Skip {vid_dir} — not found")
            continue
        vids = sorted(vid_dir.glob("*.mp4"))
        print(f"  {split}: {len(vids)} videos")
        for v in vids:
            out_dir = output_base / split / v.stem
            tasks.append((v, out_dir, split))

    print(f"\nExtracting {len(tasks)} videos with {args.workers} workers ...")

    ok = 0
    fail = 0
    with ThreadPoolExecutor(max_workers=args.workers) as ex:
        futs = {ex.submit(extract_video, t): t for t in tasks}
        from tqdm import tqdm
        for fut in tqdm(as_completed(futs), total=len(futs)):
            vid_id, n, err = fut.result()
            if err:
                print(f"  FAIL {vid_id}: {err}")
                fail += 1
            else:
                ok += 1

    print(f"\nDone: {ok} extracted, {fail} failed")
    print(f"Frames saved to: {output_base}")


if __name__ == "__main__":
    main()