#!/usr/bin/env python3
"""
run_holdout_all_models.py — replay each model's own tuned full_exp config
against the 5 held-out films, score with second_score.py, and dump a combined
JSON. r50 is excluded (see docs/model-bakeoff.md: dropped from the detailed
comparison, kept only in the calibration-curve chart).

This fills a real gap: the shipped report claimed "nothing in held-out
validation contradicts the model choice" without ever running mbf/r18 on the
held-out films — only LVFace had been checked.

Usage: python3 scripts/docs/run_holdout_all_models.py --out docs_data/holdout_all_models.json
"""
import argparse
import json
import subprocess
import sys
from pathlib import Path

REPO = Path(__file__).resolve().parent.parent.parent
sys.path.insert(0, str(REPO / "scripts" / "optimizer"))
sys.path.insert(0, str(REPO / "scripts" / "validation"))
from second_score import score_seconds  # noqa: E402
from sample_eval import load_gallery_keys  # noqa: E402

MODELS = ["LVFace-B_Glint360K", "arcface_w600k_mbf", "arcface_r18"]

HELDOUT = [
    {"name": "Benny & Joon", "slug": "Benny___Joon",
     "xray": "experiments/xray/scene_level_movie_data_XRay_US/xrays/4808_Benny__Joon"},
    {"name": "Downton Abbey: A New Era", "slug": "Downton_Abbey__A_New_Era",
     "xray": "experiments/xray/scene_level_movie_data_XRay_US/xrays/19_Downton_Abbey_A_New_Era"},
    {"name": "Lovelace", "slug": "Lovelace",
     "xray": "experiments/xray/scene_level_movie_data_XRay_US/xrays/4108_Lovelace"},
    {"name": "The Many Saints of Newark", "slug": "The_Many_Saints_of_Newark",
     "xray": "experiments/xray/scene_level_movie_data_XRay_US/xrays/900_The_Many_Saints_Of_Newark"},
    {"name": "Valerian and the City of a Thousand Planets",
     "slug": "Valerian_and_the_City_of_a_Thousand_Plan",
     "xray": "experiments/xray/scene_level_movie_data_XRay_US/xrays/5312_Valerian_and_the_City_of_a_Thousand_Planets"},
]

TRAINING = [
    {"name": "Café Society", "slug": "Café_Society",
     "xray": "experiments/xray/scene_level_movie_data_XRay_US/xrays/225_Cafe_Society"},
    {"name": "Lord of War", "slug": "Lord_of_War",
     "xray": "experiments/xray/scene_level_movie_data_XRay_US/xrays/2474_Lord_of_War"},
    {"name": "Scarface", "slug": "Scarface",
     "xray": "experiments/xray/scene_level_movie_data_XRay_US/xrays/197_Scarface"},
    {"name": "Sound of Metal", "slug": "Sound_of_Metal",
     "xray": "experiments/xray/scene_level_movie_data_XRay_US/xrays/6278_Sound_of_Metal"},
]


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--out", required=True)
    p.add_argument("--work-dir", default="/tmp/holdout_all_models")
    p.add_argument("--films", choices=["heldout", "training"], default="heldout")
    args = p.parse_args()

    work = Path(args.work_dir)
    work.mkdir(parents=True, exist_ok=True)
    film_set = HELDOUT if args.films == "heldout" else TRAINING

    results = {}
    for model in MODELS:
        cfg = json.load(open(REPO / f"experiments/results/rep4_best_{model}_full_exp.json"))["best"]["config"]
        gallery = REPO / f"experiments/galleries/gallery_{model}.h5"
        results[model] = {"config": cfg, "films": {}}

        for film in film_set:
            dump = REPO / f"experiments/dumps/{model}/dump_{film['slug']}.h5"
            if not dump.exists():
                print(f"SKIP {model}/{film['slug']}: no dump", file=sys.stderr)
                continue
            pred_path = work / f"pred_{model}_{film['slug']}.json"
            cmd = [
                "python3", "scripts/optimizer/replay.py",
                "--dump", str(dump), "--gallery", str(gallery),
                "--out", str(pred_path),
                "--prob-threshold", str(cfg["prob_threshold"]),
                "--anneal-sec", str(cfg["anneal_sec"]),
                "--extinction-sec", str(cfg["extinction_sec"]),
                "--expand-gallery",
            ]
            print(f"RUN {model}/{film['slug']}...", file=sys.stderr)
            r = subprocess.run(cmd, cwd=REPO, capture_output=True, text=True, timeout=120)
            if r.returncode != 0:
                print(f"FAIL {model}/{film['slug']}: {r.stderr[-800:]}", file=sys.stderr)
                results[model]["films"][film["slug"]] = {"error": r.stderr[-500:]}
                continue

            gk = load_gallery_keys(str(gallery))
            pred_json = json.loads(pred_path.read_text())
            m = score_seconds(pred_json, str(REPO / film["xray"]), gk)
            results[model]["films"][film["slug"]] = {"name": film["name"], **m}
            print(f"  -> F1={m['f1']*100:.1f}% P={m['precision']*100:.1f}% "
                  f"R={m['recall']*100:.1f}% misid={m['FPI_misid']}", file=sys.stderr)

    Path(args.out).parent.mkdir(parents=True, exist_ok=True)
    with open(args.out, "w") as f:
        json.dump(results, f, indent=1)
    print(f"wrote {args.out}", file=sys.stderr)


if __name__ == "__main__":
    main()