#!/usr/bin/env python3
"""Generate the 3000-video pendulum dataset + CLEVRER-style questions.

Usage:  python run_generate.py --root OUT --n-train 2400 --val-per-regime 100 --workers 2
Resume: re-running the same command skips videos/CSVs that already exist.
"""
from __future__ import annotations

import argparse
import csv
import json
import sys
import time
from multiprocessing import Pool
from pathlib import Path

import numpy as np

sys.path.insert(0, str(Path(__file__).parent))
from pendulum_gen import (FRAMES, FPS, SIZE_PX, REGIMES_VAL, REGIME_DESCRIPTIONS,
                          generate_one)
from questions import build_questions

ROOT = None


def worker(job):
    row, cfg, A = generate_one(job, str(ROOT))
    rng = np.random.default_rng(np.random.SeedSequence([job["seed"], job["index"], 777]))
    qs = build_questions(cfg, A, rng)
    entry = {
        "video_id": job["video_id"],
        "video_filename": row["video_path"],
        "split": job["split"],
        "regime": job["regime"],
        "questions": qs,
    }
    return row, entry


def init_root(root):
    global ROOT
    ROOT = root


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--root", required=True)
    ap.add_argument("--n-train", type=int, default=2400)
    ap.add_argument("--val-per-regime", type=int, default=100)
    ap.add_argument("--seed", type=int, default=20261014)
    ap.add_argument("--workers", type=int, default=2)
    args = ap.parse_args()

    root = Path(args.root)
    root.mkdir(parents=True, exist_ok=True)

    jobs = []
    for i in range(args.n_train):
        jobs.append({"seed": args.seed, "index": i, "split": "train",
                     "regime": "train", "video_id": f"video_{i:05d}"})
    k = 0
    for regime in REGIMES_VAL:
        for j in range(args.val_per_regime):
            jobs.append({"seed": args.seed, "index": 100000 + k, "split": "val",
                         "regime": regime, "video_id": f"video_{10000 + k:05d}"})
            k += 1

    print(f"{len(jobs)} videos -> {root}", flush=True)
    t0 = time.time()
    rows, q_entries = {}, {}
    done = 0
    with Pool(args.workers, initializer=init_root, initargs=(root,)) as pool:
        for row, entry in pool.imap_unordered(worker, jobs, chunksize=4):
            rows[row["video_id"]] = row
            q_entries[entry["video_id"]] = entry
            done += 1
            if done % 25 == 0 or done == len(jobs):
                rate = done / (time.time() - t0)
                eta = (len(jobs) - done) / max(rate, 1e-9)
                print(f"[{done}/{len(jobs)}] {rate:.2f} vid/s  eta {eta/60:.1f} min",
                      flush=True)

    # ---- index.csv (train first, then val; same column order as original + extras)
    cols = ["video_id", "split", "regime", "video_path", "csv_path", "frames",
            "duration_s", "fps", "size_px", "gravity_m_s2", "length_m", "name",
            "theta0_deg", "omega0_rad_s", "damping_s_inv", "damping_model",
            "quad_drag_beta", "drive_amp_rad_s2", "drive_freq_hz", "drive_phase_rad",
            "light_mean_deg", "light_amplitude_deg", "light_frequency_hz",
            "light_phase_rad"]
    order = sorted(rows, key=lambda v: (0 if rows[v]["split"] == "train" else 1, v))
    with (root / "index.csv").open("w", newline="") as f:
        w = csv.DictWriter(f, fieldnames=cols)
        w.writeheader()
        for v in order:
            w.writerow(rows[v])

    # ---- questions
    qroot = root / "questions"
    qroot.mkdir(exist_ok=True)
    for split in ("train", "val"):
        entries = [q_entries[v] for v in order if q_entries[v]["split"] == split]
        payload = {
            "info": {
                "description": "CLEVRER-style questions for the C-JEPA pendulum dataset",
                "split": split,
                "num_videos": len(entries),
                "num_questions": sum(len(e["questions"]) for e in entries),
                "question_types": ["descriptive", "explanatory", "predictive",
                                   "counterfactual"],
                "answer_convention": "descriptive: free-form string; others: index into 'choices'",
                "answers_derived_from": "ground-truth RK4 simulation (28 s; videos show first 13.33 s)",
            },
            "scene_annotation": entries,
        }
        with (qroot / f"questions_{split}.json").open("w") as f:
            json.dump(payload, f, indent=1)

    # ---- metadata.json
    meta = {
        "dataset": "pendulum_cjepa_clevrer_3k",
        "seed": args.seed,
        "generation_mode": "random within per-regime bounds",
        "num_videos": len(jobs),
        "train_videos": args.n_train,
        "val_videos": args.val_per_regime * len(REGIMES_VAL),
        "val_design": "out-of-distribution physics regimes (NOT just a held-out seed)",
        "val_regimes": REGIME_DESCRIPTIONS,
        "frames_per_video": FRAMES,
        "duration_s": FRAMES / FPS,
        "fps": FPS,
        "image_size": [SIZE_PX, SIZE_PX],
        "codec": "H.264", "pixel_format": "yuv420p", "audio": False,
        "ode": "theta_ddot = -(g/L) sin(theta) - gamma*theta_dot - beta*theta_dot*|theta_dot| + A*cos(2*pi*f_d*t+phi)",
        "integrator": "RK4", "substeps_per_frame": 8,
        "gravity_m_s2": "9.81 (train); {1.62, 3.71, 24.79} (val_gravity)",
        "length_m": 1.0,
        "split_unit": "whole video",
        "normalized_labels": False,
        "simulated_seconds_for_questions": 28.0,
        "questions_per_video_target": 8,
    }
    with (root / "metadata.json").open("w") as f:
        json.dump(meta, f, indent=2)

    n_q = sum(len(e["questions"]) for e in q_entries.values())
    print(f"DONE  {len(jobs)} videos, {n_q} questions, "
          f"{(time.time()-t0)/60:.1f} min", flush=True)


if __name__ == "__main__":
    main()