#!/usr/bin/env python3
"""Answer audit: re-derive question answers with 2x integration resolution.

The shipped answers come from the generation sim (8 substeps/frame = 96 Hz).
This script re-simulates sampled videos at 16 substeps/frame (192 Hz) and checks
that every question's answer is unchanged. Any mismatch indicates numerical or
conceptual fragility.

Usage: python audit_answers.py [dataset_root] [--n 100]
"""
import argparse
import csv
import json
import sys
from collections import Counter
from pathlib import Path

import numpy as np

sys.path.insert(0, str(Path(__file__).parent))
import pendulum_gen as G
from questions import build_questions


def resim(cfg):
    rows = G.simulate(cfg, frames=G.EXT_FRAMES, substeps=16)
    return G.rows_to_np(rows)


def cfg_from_row(r):
    cfg = G.Configuration(name=r["name"])
    cfg.theta0_deg = float(r["theta0_deg"])
    cfg.omega0_rad_s = float(r["omega0_rad_s"])
    cfg.damping_s_inv = float(r["damping_s_inv"])
    cfg.light_mean_deg = float(r["light_mean_deg"])
    cfg.light_amplitude_deg = float(r["light_amplitude_deg"])
    cfg.light_frequency_hz = float(r["light_frequency_hz"])
    cfg.light_phase_rad = float(r["light_phase_rad"])
    cfg.gravity_m_s2 = float(r["gravity_m_s2"])
    cfg.length_m = float(r["length_m"])
    cfg.damping_model = r["damping_model"]
    cfg.quad_drag_beta = float(r["quad_drag_beta"])
    cfg.drive_amp_rad_s2 = float(r["drive_amp_rad_s2"])
    cfg.drive_freq_hz = float(r["drive_freq_hz"])
    cfg.drive_phase_rad = float(r["drive_phase_rad"])
    return cfg


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("root", nargs="?", default="dataset")
    ap.add_argument("--n", type=int, default=100)
    ap.add_argument("--seed", type=int, default=20261014)
    args = ap.parse_args()

    root = Path(args.root)
    idx = list(csv.DictReader(open(root / "index.csv")))
    qs = {}
    for split in ("train", "val"):
        for e in json.load(open(root / "questions" / f"questions_{split}.json"))["scene_annotation"]:
            qs[e["video_id"]] = e["questions"]

    rng = np.random.default_rng(0)
    n_train = args.n // 2
    n_val = args.n - n_train
    pool = ([r for r in idx if r["split"] == "train"], [r for r in idx if r["split"] == "val"])
    sample = [pool[0][i] for i in rng.permutation(len(pool[0]))[:n_train]] + \
             [pool[1][i] for i in rng.permutation(len(pool[1]))[:n_val]]

    mismatches = Counter()
    n_checked = 0
    for r in sample:
        vid = r["video_id"]
        index_num = int(vid.split("_")[1]) if r["split"] == "train" else \
            100000 + (int(vid.split("_")[1]) - 10000)
        cfg = cfg_from_row(r)
        A_fine = resim(cfg)
        q_rng = np.random.default_rng(np.random.SeedSequence([args.seed, index_num, 777]))
        rebuilt = build_questions(cfg, A_fine, q_rng)
        orig = {q["question_id"]: q for q in qs[vid]}
        for q2 in rebuilt:
            q1 = orig.get(q2["question_id"])
            if q1 is None or q1["question"] != q2["question"]:
                continue
            n_checked += 1
            a1, a2 = q1["answer"], q2["answer"]
            if "choices" in q1:
                ok = q1["choices"][a1] == q2["choices"][a2]
            else:
                ok = str(a1) == str(a2)
            if not ok:
                mismatches[q2["concept"] or q2["question_type"]] += 1
                print(f"MISMATCH {vid} [{q2['question_type']}] {q2['question']}\n"
                      f"   shipped: {a1 if 'choices' not in q1 else q1['choices'][a1]}\n"
                      f"   192Hz  : {a2 if 'choices' not in q2 else q2['choices'][a2]}")
    print(f"\nchecked {n_checked} questions on {len(sample)} videos "
          f"({n_train} train / {n_val} val)")
    print(f"mismatches: {sum(mismatches.values())} {dict(mismatches)}")
    sys.exit(0 if not mismatches else 1)


if __name__ == "__main__":
    main()