#!/usr/bin/env python3
"""
run_montage_all.py — replay + best/worst-per-scene montage for all films in a
manifest, using the shipped config.hpp defaults (LVFace, prob_threshold=0.754,
anneal_sec=35.54, extinction_sec=57.43, expand_gallery). Skips a film if its
montage manifest.json already exists (safe to re-run/resume).

Manifests no longer carry a "movie" path (removed to avoid embedding source
filenames — some carry scene-release tags — in a file that gets zipped and
pushed to the artifact registry). Movie paths are resolved locally via
experiments/file-lut.json (gitignored; copy file-lut.template.json and fill
in your own paths).

Usage:
  python experiments/run_montage_all.py \
      --manifest experiments/manifests/films_LVFace-B_Glint360K_full.json \
      --gallery experiments/galleries/gallery_LVFace-B_Glint360K.h5 \
      --out-dir experiments/results/holdout/montage
"""
import argparse
import json
import subprocess
import sys
from pathlib import Path

REPO = Path(__file__).resolve().parent.parent
REPLAY = str(REPO / "scripts" / "optimizer" / "replay.py")
MONTAGE = str(REPO / "scripts" / "optimizer" / "dump_scene_montage.py")
FILE_LUT = REPO / "experiments" / "file-lut.json"

CFG = {"prob_threshold": "0.7540024664611272", "anneal_sec": "35.53996030397922",
       "extinction_sec": "57.43359645269811"}


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--manifest", required=True)
    p.add_argument("--gallery", required=True)
    p.add_argument("--out-dir", required=True)
    args = p.parse_args()

    if not FILE_LUT.exists():
        sys.exit(f"{FILE_LUT} not found — copy file-lut.template.json to "
                 f"file-lut.json and fill in your local movie paths")
    file_lut = json.loads(FILE_LUT.read_text())

    films = json.loads(Path(args.manifest).read_text())
    out_root = Path(args.out_dir)
    out_root.mkdir(parents=True, exist_ok=True)

    for f in films:
        slug = f["slug"]
        montage_dir = out_root / slug
        manifest_path = montage_dir / "manifest.json"
        if manifest_path.exists():
            print(f"[run_montage_all] skip {slug} (already done)", file=sys.stderr)
            continue

        movie_path = file_lut.get(slug)
        if not movie_path:
            print(f"[run_montage_all] skip {slug}: no entry in {FILE_LUT}", file=sys.stderr)
            continue

        print(f"[run_montage_all] === {f['name']} ({slug}) ===", file=sys.stderr)
        raw_path = out_root / f"raw_{slug}.jsonl"
        pred_path = out_root / f"pred_{slug}.json"

        replay_cmd = [sys.executable, REPLAY, "--dump", f["dump"], "--gallery", args.gallery,
                     "--out", str(pred_path), "--raw-out", str(raw_path),
                     "--prob-threshold", CFG["prob_threshold"],
                     "--anneal-sec", CFG["anneal_sec"],
                     "--extinction-sec", CFG["extinction_sec"], "--expand-gallery"]
        r = subprocess.run(replay_cmd, capture_output=True, text=True, timeout=120)
        if r.returncode != 0:
            print(f"[run_montage_all] replay FAILED for {slug}: {r.stderr[-2000:]}",
                  file=sys.stderr)
            continue

        montage_cmd = [sys.executable, MONTAGE, "--raw", str(raw_path), "--dump", f["dump"],
                       "--xray", f["xray"], "--movie", movie_path, "--gallery", args.gallery,
                       "--out-dir", str(montage_dir)]
        r = subprocess.run(montage_cmd, capture_output=True, text=True, timeout=1800)
        if r.returncode != 0:
            print(f"[run_montage_all] montage FAILED for {slug}: {r.stderr[-2000:]}",
                  file=sys.stderr)
            continue
        print(r.stderr.strip().splitlines()[-1] if r.stderr else "(no output)", file=sys.stderr)

    print("[run_montage_all] ALL DONE", file=sys.stderr)


if __name__ == "__main__":
    main()
