#!/usr/bin/env python3
"""
Build per-(model, mode) manifests for the model × gallery-mode bake-off.

Modes:
  full       — every film matches against the whole model gallery (2418 actors)
  restricted — each film matches only its Jellyfin credited cast (~15 top-billed),
               via a per-film gallery filtered from the model gallery by jellyfin_id.

Each manifest is a list of {name, xray, slug, dump, gallery} — no "movie" path (that's
resolved locally via experiments/file-lut.json, see run_montage_all.py, to avoid
embedding source filenames in a file that gets shared as an artifact). optimize.py
reads film["gallery"] per film, so restricted mode just points each film at its own
filtered gallery — no optimizer change needed.

Writes experiments/manifests/films_<model>_<mode>.json and the restricted galleries to
experiments/galleries/restricted/<model>/<slug>.h5.
"""
import json
import sys
from pathlib import Path

REPO = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(REPO / "scripts" / "validation"))
sys.path.insert(0, str(REPO / "scripts"))
from identity import norm_name  # noqa: E402
from sae_gallery import load_gallery_hdf5, save_gallery_hdf5  # noqa: E402

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

films = json.loads((REPO / "experiments/manifests/films.json").read_text())
casts = json.loads((REPO / "experiments/manifests/jellyfin_casts.json").read_text())


def restrict_gallery(model_gallery: dict, cast_ids: set[str]) -> dict:
    kept = [a for a in model_gallery["actors"] if a.get("jellyfin_id", "") in cast_ids]
    return {"actors": kept}


for model in MODELS:
    gpath = REPO / f"experiments/galleries/gallery_{model}.h5"
    if not gpath.exists():
        print(f"skip {model}: gallery not built yet ({gpath})")
        continue
    model_gal = load_gallery_hdf5(gpath)

    # full mode
    full = [{**f, "dump": f"experiments/dumps/{model}/dump_{f['slug']}.h5",
             "gallery": f"experiments/galleries/gallery_{model}.h5"} for f in films]
    (REPO / f"experiments/manifests/films_{model}_full.json").write_text(json.dumps(full, indent=2, ensure_ascii=False))

    # restricted mode — per-film filtered gallery
    rdir = REPO / f"experiments/galleries/restricted/{model}"
    rdir.mkdir(parents=True, exist_ok=True)
    restr = []
    for f in films:
        cast_ids = set(casts.get(f["name"], []))
        rg = restrict_gallery(model_gal, cast_ids)
        rgpath = rdir / f"{f['slug']}.h5"
        save_gallery_hdf5(rg, rgpath)
        restr.append({**f, "dump": f"experiments/dumps/{model}/dump_{f['slug']}.h5",
                      "gallery": str(rgpath.relative_to(REPO)),
                      "_cast_size": len(rg["actors"])})
    (REPO / f"experiments/manifests/films_{model}_restricted.json").write_text(json.dumps(restr, indent=2, ensure_ascii=False))
    avg = sum(r["_cast_size"] for r in restr) / len(restr)
    print(f"{model}: full (2418) + restricted (avg {avg:.0f} actors/film) manifests written")