#!/usr/bin/env python3
"""
min_face_size.py — VR-005: at what face size do embeddings stop identifying people?

TRACES: VR-005

`min_face_px` is currently a working estimate (AR-002: 66x66 px in original video
resolution). This script replaces the guess with a measurement, using only gallery
mugshots already on disk — no video, no C++ changes.

Protocol
--------
1. Select ~100 gallery actors that have more than one mugshot.
2. Per actor hold out ONE image as the *probe*; that actor's remaining images stay
   in the gallery at native resolution.
3. For each target size S, take the probe's native aligned 112x112 crop, downscale
   it to SxS and upscale it back to 112x112, then embed. Detail is genuinely
   destroyed and then the same warp the pipeline applies is re-applied on top —
   which is what a face detected at SxS in a frame actually suffers.
4. Match each degraded probe against the whole gallery.
5. Record, per size, TPI (identified as the correct actor) and FPI (identified as
   someone else). Everything else is an unidentified probe (TBI).

The asymmetry is the point: **the gallery stays at native resolution and only the
probe degrades.** That is the production case — reference mugshots are clean, the
face coming out of the video is small. Degrading both sides would measure
something the pipeline never does.

What this deliberately does NOT measure
---------------------------------------
The cosine between the size-S embedding and the native embedding of the *same*
image. That is embedding *drift*, and it answers the wrong question: an embedding
can drift a long way and stay perfectly separable, or drift a little in a
direction that destroys separation. What matters is the decision the pipeline
makes — probe against a competing gallery — so that is what is recorded.

CAVEAT — FPI IS RELATIVE, NOT ABSOLUTE
--------------------------------------
False positives grow with the number of actors competing for the match. A
~100-actor gallery therefore *understates* the false-positive rate against a
production library of thousands. Read the FPI column as a relative curve across
sizes ("FPI is 4x worse at 32 px than at 64 px"), never as the rate you would see
in production. Re-run with `--actors` at production scale before setting a
threshold from an absolute FPI number.

Decision rule
-------------
Per the repo invariant (CLAUDE.md: "always use the calibrated probability, never a
raw cosine"), identification goes through the same path as `identity_matcher_node`:
per-actor best-of-N cosine -> Platt sigmoid P(match) = sigma(a*sim + b + log-prior)
-> accept if P > `prob_threshold`. The sigmoid is fitted here by the same
histogram/gradient-descent procedure as `src/gallery/gallery_calibration.hpp`,
over the native gallery embeddings only (held-out probes are excluded, so the
calibration cannot see the images it will be scored on).

How this runs
-------------
Through the `sae_embed` bindings, which expose the shipped C++ stages directly:
`detect()`, `align_face()`, `embed_crops()` and `GalleryCalibration`. Nothing
here re-implements detection, the ArcFace warp, the embedder or the Platt fit.

That matters most for the calibration. A second copy of the sigmoid is exactly
where "always the calibrated probability, never a raw cosine" (AR-024) gets
broken without anyone noticing, because the copy keeps returning plausible
numbers after the original has moved. Scoring through the binding makes the rule
structural instead of remembered.

The backend is whichever was compiled in. Under `SAE_INFERENCE_BACKEND=ORT`
that is the reference fp32 path, which loads the .onnx directly. A TensorRT fp16
build is a *different realisation* of the same model and its embeddings are
measurably not the same vectors: on LVFace-B_Glint360K the stored TRT-fp16 gallery
agrees with an fp32 recompute of the same mugshot at only ~0.85 cosine, while
same-actor/different-actor separation is essentially unchanged (d' 5.3 vs 5.7).
Nothing here is invalidated by that — gallery and probes go through one session,
so the comparison is internally consistent — but the two embedding spaces are not
interchangeable, and `--verify-against <gallery.h5>` will show ~0.85, not ~1.0,
against a TRT-built gallery. It reports the separation of both sets alongside the
agreement so the two causes are distinguishable: a broken port collapses
separation, a different backend does not.

Secondary output (VR-005): running the sweep per `--arcface` model shows whether
`min_face_px` should be one constant at all, or should scale with the embedder —
which matters because the model is a build-time choice (GR-004).

Usage
-----
  python scripts/validation/min_face_size.py \
      --images images \
      --gallery gallery_lvface.h5 \
      --arcface models/LVFace-B_Glint360K.onnx \
      --actors 100 --seed 0 \
      --out experiments/results/vr005_min_face_size

Writes <out>.csv, <out>.json and <out>.png (plus <out>.per_probe.csv with
--per-probe).
"""
from __future__ import annotations

import argparse
import csv
import json
import random
import re
import sys
import time
from pathlib import Path

REPO = Path(__file__).resolve().parent.parent.parent
sys.path.insert(0, str(REPO / "scripts"))
def _find_sae_embed() -> Path | None:
    """Locate the built sae_embed module.

    A git worktree has no build tree of its own, so fall back to the main
    checkout via the shared git dir — otherwise running this study from a
    feature worktree cannot find the bindings it now depends on.
    """
    roots = [REPO]
    try:
        import subprocess
        common = subprocess.run(["git", "-C", str(REPO), "rev-parse",
                                 "--path-format=absolute", "--git-common-dir"],
                                capture_output=True, text=True, check=True).stdout.strip()
        if common:
            roots.append(Path(common).parent)
    except Exception:
        pass
    for root in roots:
        for b in ("build-ort", "build"):
            if list((root / b).glob("sae_embed*.so")):
                return root / b
    return None


_SAE_BUILD = _find_sae_embed()
if _SAE_BUILD is None:
    sys.exit("cannot find the built sae_embed module — build it with\n"
             "  cmake --build build-ort --target sae_embed")
sys.path.insert(0, str(_SAE_BUILD))

# Before cv2: OpenCV's DNN module loads the system libonnxruntime, which then
# shadows the one sae_embed links against and the import fails on a missing
# symbol version. Order matters here.
import sae_embed

import cv2
import numpy as np

IMAGE_EXTS = (".jpg", ".jpeg", ".png", ".webp")
JELLYFIN_ID_RE = re.compile(r"^[0-9a-f]{32}$")

# Interpolation used for the two halves of the degradation. Downscaling uses
# INTER_AREA (correct low-pass for shrinking, i.e. detail that a small detection
# genuinely never had); upscaling uses INTER_LINEAR, which is what warpAffine in
# align_face() uses when it blows a small detection up to 112x112.
INTERP = {
    "area":    cv2.INTER_AREA,
    "linear":  cv2.INTER_LINEAR,
    "cubic":   cv2.INTER_CUBIC,
    "nearest": cv2.INTER_NEAREST,
    "lanczos": cv2.INTER_LANCZOS4,
}

# House chart palette, shared with scripts/docs/experiment_charts.py so figures
# across the report read as one set.
INK, MUTED, GRID, SURFACE = "#0b0b0b", "#898781", "#e1e0d9", "#fcfcfb"
BLUE, GREEN, RED, AMBER = "#2a78d6", "#008300", "#e34948", "#eda100"


# ── Production stages, via the sae_embed bindings ─────────────────────────────
# detect / align_face / embed_crops / calibrate_gallery all call the shipped C++.
# There is deliberately no Python re-implementation of any of them: a second copy
# drifts from what ships, and the calibration is the one that must not — AR-024
# requires every similarity to pass through the same sigmoid the matcher uses.

# These two mirror constants in gallery_calibration.hpp. They are NOT a second
# copy of the fit — that is the binding's job — but the script reproduces the
# same dedup and eligibility filtering so the actor counts it reports describe
# the population the C++ actually fitted on. Keep them in step with the header.
MIN_EMB_FOR_POSITIVE = 5
DEDUP_SIM = 1.0 - 1e-7


class Stages:
    """Thin holder so the rest of the script has one object to call."""

    def __init__(self, detector: str, arcface: str, conf: float, nms: float):
        self.engine = sae_embed.FaceEmbedder(
            detector_model=detector, arcface_model=arcface,
            conf=conf, nms=nms, max_side=0)

    def detect(self, img):
        return self.engine.detect(img)

    def align(self, img, landmarks):
        return sae_embed.align_face(img, np.asarray(landmarks, dtype=np.float32).reshape(5, 2))

    def enhance(self, img):
        return sae_embed.enhance_for_retry(img)

    def embed(self, crops):
        """(N,112,112,3) uint8 BGR -> (N,512) float32.

        Chunked at the backend's max_batch: the engine does not split an
        oversized request, so handing it a whole gallery at once asks CUDA for
        a multi-gigabyte activation buffer and the allocator refuses.
        """
        if not len(crops):
            return np.zeros((0, 512), dtype=np.float32)
        n = max(1, int(self.engine.max_batch))
        arr = np.ascontiguousarray(np.stack(crops), dtype=np.uint8)
        out = [np.asarray(self.engine.embed_crops(np.ascontiguousarray(arr[i:i + n])))
               for i in range(0, len(arr), n)]
        return np.concatenate(out, axis=0)


def calibrate_gallery(emb: np.ndarray, actor: np.ndarray) -> dict:
    """The production Platt fit (gallery_calibration.hpp), via the binding."""
    cal = sae_embed.calibrate_gallery(
        np.ascontiguousarray(emb, dtype=np.float32), [int(a) for a in actor])
    print(f"[calibration] a={cal.a:.4f} b={cal.b:.4f} valid={cal.valid} "
          f"boundary(P=0.5)=sim{cal.boundary_at(0.5):.4f}", file=sys.stderr)
    # Held module-side rather than returned: the returned dict lands in the run
    # metadata, and a native object there breaks the JSON dump.
    _CAL["cal"] = cal
    return {"a": float(cal.a), "b": float(cal.b), "valid": bool(cal.valid)}


def probability(sim, a: float, b: float, log_prior_odds: float = 0.0):
    """P(match) through GalleryCalibration — the C++ sigmoid, not a copy of it."""
    cal = _CAL.get("cal")
    if cal is None:
        raise RuntimeError("probability() called before calibrate_gallery()")
    sim = np.asarray(sim, dtype=np.float64)
    flat = np.atleast_1d(sim).ravel()
    out = np.array([cal.probability(float(v), log_prior_odds) for v in flat])
    return out.reshape(sim.shape) if sim.shape else float(out[0])


_CAL: dict = {}


# ── Runtime / actor discovery ─────────────────────────────────────────────────

def normalise_name(name: str) -> str:
    return re.sub(r"[^a-z0-9]+", "", name.lower())


def discover_actors(images_root: Path) -> list[dict]:
    """Enumerate the gallery-build image cache: <root>/<jellyfin_id>_<Name>/NN.jpg
    (the layout make_jellyfin_gallery.py / reembed_gallery.py use)."""
    actors = []
    for d in sorted(p for p in images_root.iterdir() if p.is_dir()):
        imgs = sorted(p for p in d.iterdir()
                      if p.is_file() and p.suffix.lower() in IMAGE_EXTS)
        if not imgs:
            continue
        head, _, tail = d.name.partition("_")
        if JELLYFIN_ID_RE.match(head) and tail:
            jellyfin_id, name = head, tail.replace("_", " ")
        else:
            jellyfin_id, name = "", d.name.replace("_", " ")
        actors.append({"dir": d, "jellyfin_id": jellyfin_id, "name": name,
                       "images": imgs})
    return actors


def gallery_keys(gallery_path: Path) -> tuple[set[str], set[str]]:
    """(jellyfin ids, normalised names) of the actors an existing gallery holds."""
    from sae_gallery import load_gallery_hdf5
    g = load_gallery_hdf5(gallery_path)
    ids = {a.get("jellyfin_id", "") for a in g["actors"] if a.get("jellyfin_id")}
    names = {normalise_name(a.get("name", "")) for a in g["actors"] if a.get("name")}
    return ids, names


# ── Degradation ───────────────────────────────────────────────────────────────

def degrade(crop: np.ndarray, size: int, down: int, up: int) -> np.ndarray:
    """Throw away everything a face detected at size x size never had, then warp
    it back up to the 112x112 the embedder is fed."""
    if size == 112:
        return crop
    small = cv2.resize(crop, (size, size), interpolation=down)
    return cv2.resize(small, (112, 112), interpolation=up)


# ── Reporting ─────────────────────────────────────────────────────────────────

CAVEAT = (
    "CAVEAT: FPI grows with gallery size. This ran against {n_actors} actors, so it "
    "UNDERSTATES the false-positive rate of a production library of thousands. Read "
    "FPI as relative across sizes, not as an absolute rate."
)


def write_plot(rows: list[dict], out_png: Path, meta: dict) -> None:
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    plt.rcParams.update({
        "figure.facecolor": SURFACE, "axes.facecolor": SURFACE,
        "savefig.facecolor": SURFACE, "text.color": INK,
        "axes.edgecolor": MUTED, "axes.labelcolor": INK,
        "xtick.color": MUTED, "ytick.color": MUTED,
        "axes.grid": True, "grid.color": GRID, "grid.linewidth": 0.8,
        "axes.spines.top": False, "axes.spines.right": False,
    })

    sizes = [r["size_px"] for r in rows]
    fig, ax = plt.subplots(figsize=(9, 5.6))
    ax.plot(sizes, [100 * r["tpi_rate"] for r in rows], "-o", color=GREEN,
            lw=2, label="TPI — identified, correct actor")
    ax.plot(sizes, [100 * r["fpi_rate"] for r in rows], "-s", color=RED,
            lw=2, label="FPI — identified, wrong actor")
    ax.plot(sizes, [100 * r["unidentified_rate"] for r in rows], color=MUTED,
            marker="^", lw=1.4, ls="--", label="unidentified (below P threshold)")
    ax.plot(sizes, [100 * r["rank1_rate"] for r in rows], ":", color=BLUE,
            lw=1.6, label="rank-1 correct (ignoring threshold)")

    op = meta.get("operating_point")
    if op:
        ax.axvline(op, color=AMBER, lw=1.6, ls="-.", zorder=1)
        ax.annotate(f"operating point {op} px", xy=(op, 50),
                    xytext=(4, 0), textcoords="offset points",
                    color=AMBER, fontsize=9, rotation=90, va="center")

    ax.set_xlabel("probe face size before upscaling (px)")
    ax.set_ylabel("% of probes")
    ax.set_ylim(-2, 102)
    ax.set_xticks(sizes)
    ax.set_title(f"VR-005 — identification vs. probe face size\n"
                 f"{meta['model']}, {meta['n_actors']} actors, "
                 f"{meta['n_probes']} probes/size, gallery at native resolution",
                 fontsize=11, loc="left")
    ax.legend(frameon=False, fontsize=9, loc="center left")
    fig.text(0.01, 0.005, CAVEAT.format(n_actors=meta["n_actors"]),
             fontsize=7.5, color=MUTED, wrap=True)
    fig.tight_layout(rect=(0, 0.05, 1, 1))
    out_png.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(out_png, dpi=150)
    plt.close(fig)


def pick_operating_point(rows: list[dict], retention: float, fpi_slack: float) -> int | None:
    """Smallest size that keeps `retention` of the undegraded (112 px control)
    TPI rate and does not add more than `fpi_slack` absolute FPI over it.
    A stated rule, not a magic number — change the rule, not the answer."""
    control = next((r for r in rows if r["size_px"] == 112), None)
    if control is None or control["n_probes"] == 0:
        return None
    tpi_floor = retention * control["tpi_rate"]
    fpi_ceil = control["fpi_rate"] + fpi_slack
    ok = [r["size_px"] for r in rows
          if r["tpi_rate"] >= tpi_floor and r["fpi_rate"] <= fpi_ceil]
    return min(ok) if ok else None


# ── Main ──────────────────────────────────────────────────────────────────────

def main() -> int:
    p = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("--images", required=True,
                   help="gallery image cache root (<jellyfin_id>_<Name>/NN.jpg)")
    p.add_argument("--gallery", default=None,
                   help="gallery .h5 — restricts the actor pool to its members")
    p.add_argument("--out", default=str(REPO / "experiments/results/vr005_min_face_size"),
                   help="output path prefix (.csv/.json/.png are appended)")
    p.add_argument("--per-probe", action="store_true",
                   help="also write <out>.per_probe.csv, one row per probe per size")

    p.add_argument("--actors", type=int, default=100, help="actors to sample (default 100)")
    p.add_argument("--min-images", type=int, default=2,
                   help="minimum mugshots for an actor to be eligible (default 2)")
    p.add_argument("--probes-per-actor", type=int, default=1,
                   help="images held out per actor; 1 is the VR-005 protocol")
    p.add_argument("--seed", type=int, default=0, help="actor/probe selection seed")
    p.add_argument("--keep-duplicates", action="store_true",
                   help="keep mugshots that are the same photograph twice; by "
                        "default they are dropped, since a probe identical to a "
                        "gallery reference is identified for free at every size")
    p.add_argument("--sizes", default="12,16,20,24,32,40,48,56,64,72,80,96,112",
                   help="comma-separated probe sizes; 112 is the undegraded control")

    p.add_argument("--models-dir", default=str(REPO / "models"))
    p.add_argument("--arcface", default=None,
                   help="embedder ONNX (default <models-dir>/LVFace-B_Glint360K.onnx)")
    p.add_argument("--detector", default=None,
                   help="SCRFD ONNX (default <models-dir>/scrfd_500m_bnkps.onnx)")
    p.add_argument("--conf", type=float, default=0.5, help="detector confidence")
    p.add_argument("--nms", type=float, default=0.4, help="detector NMS IoU")
    p.add_argument("--max-side", type=int, default=500,
                   help="downscale mugshots to this longest side before detection, "
                        "matching the gallery builders' embedder settings")

    p.add_argument("--prob-threshold", type=float, default=0.754,
                   help="accept if P(match) exceeds this (Config::prob_threshold)")
    p.add_argument("--match-prior", type=float, default=0.5,
                   help="base-rate prior (Config::match_prior)")
    p.add_argument("--calib-a", type=float, default=None,
                   help="override the fitted sigmoid scale instead of fitting")
    p.add_argument("--calib-b", type=float, default=None,
                   help="override the fitted sigmoid bias instead of fitting")

    p.add_argument("--down-interp", default="area", choices=sorted(INTERP),
                   help="interpolation for the 112 -> S downscale (default area)")
    p.add_argument("--up-interp", default="linear", choices=sorted(INTERP),
                   help="interpolation for the S -> 112 upscale (default linear, "
                        "as warpAffine uses in align_face)")

    p.add_argument("--tpi-retention", type=float, default=0.95,
                   help="operating point keeps this fraction of the control TPI rate")
    p.add_argument("--fpi-slack", type=float, default=0.01,
                   help="operating point may add at most this absolute FPI over control")
    p.add_argument("--verify-against", default=None,
                   help="gallery .h5 built from --images with the same model: report "
                        "agreement and separation of recomputed vs stored embeddings "
                        "(a TensorRT-built gallery will not agree; see the docstring)")
    args = p.parse_args()

    if (args.calib_a is None) != (args.calib_b is None):
        return err("--calib-a and --calib-b must be given together")

    models_dir = Path(args.models_dir)
    arcface = Path(args.arcface) if args.arcface else models_dir / "LVFace-B_Glint360K.onnx"
    detector = Path(args.detector) if args.detector else models_dir / "scrfd_500m_bnkps.onnx"
    for path, what in ((arcface, "embedder"), (detector, "detector")):
        if not path.is_file():
            return err(f"{what} model not found: {path}\n"
                       f"Run: bash scripts/download_models.sh")

    images_root = Path(args.images)
    if not images_root.is_dir():
        return err(f"image cache not found: {images_root}")

    sizes = sorted({int(s) for s in args.sizes.split(",") if s.strip()})
    if not sizes:
        return err("--sizes is empty")
    if args.probes_per_actor < 1:
        return err("--probes-per-actor must be >= 1")

    cv2.setRNGSeed(args.seed)  # estimateAffinePartial2D's RANSAC draws from this

    # ── actor pool ────────────────────────────────────────────────────────────
    pool = discover_actors(images_root)
    print(f"[select] {len(pool)} actor dirs with images under {images_root}",
          file=sys.stderr)
    if args.gallery:
        ids, names = gallery_keys(Path(args.gallery))
        pool = [a for a in pool
                if (a["jellyfin_id"] and a["jellyfin_id"] in ids)
                or normalise_name(a["name"]) in names]
        print(f"[select] {len(pool)} of them are in {args.gallery}", file=sys.stderr)

    need = max(args.min_images, args.probes_per_actor + 1)
    eligible = [a for a in pool if len(a["images"]) >= need]
    print(f"[select] {len(eligible)} have >= {need} mugshots", file=sys.stderr)
    if len(eligible) < 2:
        return err(f"need at least 2 actors with >= {need} mugshots; found "
                   f"{len(eligible)}. Build the image cache first "
                   f"(scripts/make_jellyfin_gallery.py) or lower --min-images.")

    rng = random.Random(args.seed)
    selected = sorted(rng.sample(eligible, min(args.actors, len(eligible))),
                      key=lambda a: a["dir"].name)
    if len(selected) < args.actors:
        print(f"[select] WARNING: only {len(selected)} eligible actors, "
              f"--actors {args.actors} requested. FPI is gallery-size dependent — "
              f"see the caveat.", file=sys.stderr)

    # ── detect + align every mugshot of the selected actors, once ─────────────
    stages = Stages(str(detector), str(arcface), args.conf, args.nms)
    print(f"[models] detector={detector.name} embedder={arcface.name} "
          f"batch={stages.engine.max_batch} (provider chosen by the C++ backend: "
          f"CUDA, then ROCm, then CPU)", file=sys.stderr)

    t0 = time.time()
    crops: list[np.ndarray] = []
    rows: list[dict] = []          # parallel to crops: {actor, actor_idx, image}
    actors: list[dict] = []
    n_nodetect = 0
    for a in selected:
        actor_crops, actor_paths = [], []
        for img_path in a["images"]:
            img = cv2.imread(str(img_path))
            if img is None:
                n_nodetect += 1
                continue
            if args.max_side > 0 and max(img.shape[:2]) > args.max_side:
                s = args.max_side / max(img.shape[:2])
                img = cv2.resize(img, None, fx=s, fy=s, interpolation=cv2.INTER_AREA)
            faces = stages.detect(img)
            if not faces:
                enhanced = stages.enhance(img)
                faces = stages.detect(enhanced)
                if faces:
                    img = enhanced
            if not faces:
                n_nodetect += 1
                continue
            best = max(faces, key=lambda f: f.confidence)
            crop = stages.align(img, best.landmarks)
            if crop is None:
                n_nodetect += 1
                continue
            actor_crops.append(crop)
            actor_paths.append(img_path)
        if len(actor_crops) < args.probes_per_actor + 1:
            continue
        ai = len(actors)
        actors.append({"name": a["name"], "jellyfin_id": a["jellyfin_id"],
                       "dir": a["dir"].name, "n_images": len(actor_crops)})
        for crop, img_path in zip(actor_crops, actor_paths):
            rows.append({"actor_idx": ai, "image": str(img_path)})
            crops.append(crop)
        if len(actors) % 20 == 0:
            print(f"  [align] {len(actors)}/{len(selected)} actors, "
                  f"{len(crops)} crops", file=sys.stderr)

    if len(actors) < 2:
        return err(f"only {len(actors)} actors survived detection/alignment — "
                   f"nothing to match against")
    print(f"[align] {len(actors)} actors, {len(crops)} aligned crops, "
          f"{n_nodetect} images skipped (no face / unreadable) in "
          f"{time.time() - t0:.1f}s", file=sys.stderr)

    actor_of = np.array([r["actor_idx"] for r in rows], dtype=int)

    # ── embed everything at native resolution ────────────────────────────────
    t0 = time.time()
    native = stages.embed(crops)
    print(f"[embed] {len(crops)} native crops in {time.time() - t0:.1f}s",
          file=sys.stderr)

    if args.verify_against:
        verify_embeddings(Path(args.verify_against), rows, native, actor_of)

    # ── drop duplicate mugshots ───────────────────────────────────────────────
    # The cache holds the same photograph twice for some actors (two provider
    # URLs, one picture). A probe that is identical to a gallery reference is
    # identified for free at every size, which flatters the whole curve, so
    # remove duplicates the same way calibrate_gallery does.
    n_dup = 0
    if not args.keep_duplicates:
        keep = np.ones(len(rows), bool)
        for ai in range(len(actors)):
            kept: list[int] = []
            for i in np.nonzero(actor_of == ai)[0]:
                if any(float(native[i] @ native[k]) > DEDUP_SIM for k in kept):
                    keep[i] = False
                else:
                    kept.append(int(i))
        n_dup = int((~keep).sum())

        # An actor left with too few distinct mugshots to hold one out drops out.
        counts = np.bincount(actor_of[keep], minlength=len(actors))
        drop_actor = counts < args.probes_per_actor + 1
        keep &= ~drop_actor[actor_of]

        remap = np.full(len(actors), -1, dtype=int)
        remap[~drop_actor] = np.arange(int((~drop_actor).sum()))
        actors = [a for a, d in zip(actors, drop_actor) if not d]
        rows = [r for r, k in zip(rows, keep) if k]
        crops = [c for c, k in zip(crops, keep) if k]
        native = native[keep]
        actor_of = remap[actor_of[keep]]
        for r, ai in zip(rows, actor_of):
            r["actor_idx"] = int(ai)
        for ai, a in enumerate(actors):
            a["n_images"] = int(np.sum(actor_of == ai))
        print(f"[dedup] dropped {n_dup} duplicate mugshots and "
              f"{int(drop_actor.sum())} actors left with too few; "
              f"{len(actors)} actors, {len(rows)} images remain", file=sys.stderr)
        if len(actors) < 2:
            return err("fewer than 2 actors survive de-duplication — the image "
                       "cache holds too few distinct mugshots")

    # ── hold out the probes ───────────────────────────────────────────────────
    is_probe = np.zeros(len(rows), bool)
    for ai in range(len(actors)):
        idx = np.nonzero(actor_of == ai)[0]
        # Seeded per actor so the choice does not depend on iteration order.
        r = random.Random(f"{args.seed}:{actors[ai]['dir']}")
        for pick in r.sample(list(idx), args.probes_per_actor):
            is_probe[pick] = True
    probe_rows = np.nonzero(is_probe)[0]
    gal_rows = np.nonzero(~is_probe)[0]
    print(f"[holdout] {len(probe_rows)} probes held out, "
          f"{len(gal_rows)} gallery embeddings remain", file=sys.stderr)

    gal_emb = native[gal_rows]
    gal_actor = actor_of[gal_rows]
    probe_actor = actor_of[probe_rows]

    # Per-actor column masks for the best-of-N scan (identity_matcher_node).
    actor_cols = [np.nonzero(gal_actor == ai)[0] for ai in range(len(actors))]
    have_refs = np.array([len(c) > 0 for c in actor_cols])
    if not have_refs.all():
        return err("an actor ended up with no gallery references left; "
                   "raise --min-images")

    # ── calibration ───────────────────────────────────────────────────────────
    if args.calib_a is not None:
        cal = {"a": args.calib_a, "b": args.calib_b, "valid": True}
        print(f"[calibration] using supplied a={cal['a']} b={cal['b']}", file=sys.stderr)
    else:
        cal = calibrate_gallery(gal_emb, gal_actor)
    if not cal["valid"]:
        return err(
            "calibration could not be fitted, and this study will not fall back to a "
            "raw cosine threshold (CLAUDE.md invariant). Use more actors with >= "
            f"{MIN_EMB_FOR_POSITIVE} mugshots, or pass --calib-a/--calib-b from a "
            "production gallery.")
    log_prior_odds = float(np.log(args.match_prior / (1.0 - args.match_prior)))

    # ── sweep ─────────────────────────────────────────────────────────────────
    down, up = INTERP[args.down_interp], INTERP[args.up_interp]
    probe_crops = [crops[i] for i in probe_rows]
    results, per_probe = [], []
    for size in sizes:
        t0 = time.time()
        degraded = [degrade(c, size, down, up) for c in probe_crops]
        q = stages.embed(degraded)

        sims = q @ gal_emb.T                                   # [n_probe, n_gal]
        best_per_actor = np.stack([sims[:, cols].max(axis=1) for cols in actor_cols],
                                  axis=1)                       # [n_probe, n_actor]
        best_actor = best_per_actor.argmax(axis=1)
        best_sim = best_per_actor.max(axis=1)
        p_match = np.asarray(probability(best_sim, cal["a"], cal["b"], log_prior_odds))

        accept = p_match > args.prob_threshold
        correct = best_actor == probe_actor
        tpi = int(np.sum(accept & correct))
        fpi = int(np.sum(accept & ~correct))
        unid = int(np.sum(~accept))
        n = len(probe_rows)

        results.append({
            "size_px": size,
            "n_probes": n,
            "tpi": tpi, "fpi": fpi, "unidentified": unid,
            "tpi_rate": tpi / n, "fpi_rate": fpi / n, "unidentified_rate": unid / n,
            "rank1_rate": float(np.mean(correct)),
            "mean_best_sim": float(np.mean(best_sim)),
            "mean_p_match": float(np.mean(p_match)),
            "mean_sim_true_actor": float(np.mean(
                best_per_actor[np.arange(n), probe_actor])),
        })
        if args.per_probe:
            for j in range(n):
                per_probe.append({
                    "size_px": size,
                    "probe_image": rows[probe_rows[j]]["image"],
                    "true_actor": actors[probe_actor[j]]["name"],
                    "matched_actor": actors[best_actor[j]]["name"],
                    "best_sim": float(best_sim[j]),
                    "p_match": float(p_match[j]),
                    "outcome": ("TPI" if accept[j] and correct[j]
                                else "FPI" if accept[j] else "unidentified"),
                })
        print(f"[sweep] {size:3d}px  TPI {tpi:4d} ({100 * tpi / n:5.1f}%)  "
              f"FPI {fpi:4d} ({100 * fpi / n:5.1f}%)  "
              f"unid {unid:4d} ({100 * unid / n:5.1f}%)  "
              f"rank1 {100 * np.mean(correct):5.1f}%  "
              f"[{time.time() - t0:.1f}s]", file=sys.stderr)

    op = pick_operating_point(results, args.tpi_retention, args.fpi_slack)

    # ── outputs ───────────────────────────────────────────────────────────────
    out = Path(args.out)
    out.parent.mkdir(parents=True, exist_ok=True)
    # Append rather than with_suffix() so a prefix containing a dot keeps its name.
    csv_path = out.with_name(out.name + ".csv")
    json_path = out.with_name(out.name + ".json")
    png_path = out.with_name(out.name + ".png")
    fields = list(results[0].keys())
    with open(csv_path, "w", newline="") as f:
        w = csv.DictWriter(f, fieldnames=fields)
        w.writeheader()
        w.writerows(results)

    meta = {
        "requirement": "VR-005",
        "caveat": CAVEAT.format(n_actors=len(actors)),
        "model": arcface.stem,
        "detector": detector.stem,
        "backend": "sae_embed / the compiled-in inference backend (fp32 ONNX under "
                   "SAE_INFERENCE_BACKEND=ORT; a TensorRT fp16 build is a different "
                   "embedding space)",
        "n_actors": len(actors),
        "n_probes": len(probe_rows),
        "n_gallery_embeddings": len(gal_rows),
        "probes_per_actor": args.probes_per_actor,
        "seed": args.seed,
        "sizes": sizes,
        "prob_threshold": args.prob_threshold,
        "match_prior": args.match_prior,
        "calibration": cal,
        "calibration_source": "supplied" if args.calib_a is not None else "fitted",
        "sim_boundary_at_threshold": float(
            (np.log(args.prob_threshold / (1 - args.prob_threshold))
             - cal["b"] - log_prior_odds) / cal["a"]),
        "down_interp": args.down_interp,
        "up_interp": args.up_interp,
        "max_side": args.max_side,
        "duplicate_mugshots_dropped": n_dup,
        "operating_point_rule": (
            f"smallest size retaining >= {args.tpi_retention:.0%} of the 112 px "
            f"control TPI rate with <= +{args.fpi_slack:.1%} absolute FPI"),
        "operating_point": op,
        "images_skipped_no_face": n_nodetect,
        "curve": results,
        "actors": actors,
    }
    json_path.write_text(json.dumps(meta, indent=2) + "\n")

    if per_probe:
        pp = out.with_name(out.name + ".per_probe.csv")
        with open(pp, "w", newline="") as f:
            w = csv.DictWriter(f, fieldnames=list(per_probe[0].keys()))
            w.writeheader()
            w.writerows(per_probe)
        print(f"[out] {pp}", file=sys.stderr)

    write_plot(results, png_path, meta)

    # ── stdout report ─────────────────────────────────────────────────────────
    print(f"\nVR-005 — minimum face size, {arcface.stem}")
    print(f"{len(actors)} actors, {len(probe_rows)} probes/size, "
          f"{len(gal_rows)} gallery embeddings at native resolution")
    print(f"identify when P>{args.prob_threshold}, i.e. cosine above "
          f"{meta['sim_boundary_at_threshold']:.4f} under the calibration fitted "
          f"on this gallery\n")
    print(f"{'size':>5} {'TPI':>8} {'FPI':>8} {'unid':>8} {'rank1':>8} {'mean sim':>9}")
    for r in results:
        print(f"{r['size_px']:>5} {100 * r['tpi_rate']:>7.1f}% "
              f"{100 * r['fpi_rate']:>7.1f}% {100 * r['unidentified_rate']:>7.1f}% "
              f"{100 * r['rank1_rate']:>7.1f}% {r['mean_best_sim']:>9.4f}")
    print(f"\noperating point: {op if op else 'none of the swept sizes qualifies'}"
          f"   ({meta['operating_point_rule']})")
    print(f"\n{meta['caveat']}")
    print(f"\n[out] {csv_path}\n[out] {json_path}\n[out] {png_path}")
    return 0


def _separation(emb: np.ndarray, actor: np.ndarray) -> tuple[float, float, float]:
    """(mean same-actor sim, mean different-actor sim, d') — the property that has
    to survive for an embedding space to be usable, whatever its coordinates."""
    iu, ju = np.triu_indices(len(actor), k=1)
    sims = (emb @ emb.T)[iu, ju]
    same = actor[iu] == actor[ju]
    pos = sims[same & (sims < 0.9999)]   # drop duplicate source images
    neg = sims[~same]
    if pos.size < 2 or neg.size < 2:
        return float("nan"), float("nan"), float("nan")
    d = (pos.mean() - neg.mean()) / np.sqrt(0.5 * (pos.var() + neg.var()))
    return float(pos.mean()), float(neg.mean()), float(d)


def verify_embeddings(gallery_path: Path, rows: list[dict], native: np.ndarray,
                      actor_of: np.ndarray) -> None:
    """Cross-check this script's ONNX port against a gallery built by the C++
    pipeline from the same mugshots.

    Agreement is ~1.0 only if that gallery was built with the same backend. A
    TensorRT fp16 build lands around 0.85 on LVFace-B while separating just as
    well, so the separation figures — not the agreement — are what says whether
    the port is sound."""
    from sae_gallery import load_gallery_hdf5
    g = load_gallery_hdf5(gallery_path)
    stored: dict[tuple[str, str], np.ndarray] = {}
    for a in g["actors"]:
        key = a.get("jellyfin_id") or normalise_name(a.get("name", ""))
        for e, src in zip(a.get("embeddings", []), a.get("source_images", [])):
            if src:
                stored[(key, src)] = np.asarray(e, np.float32)

    sims, paired_mine, paired_ref, paired_actor = [], [], [], []
    for i, r in enumerate(rows):
        path = Path(r["image"])
        head, _, _ = path.parent.name.partition("_")
        key = head if JELLYFIN_ID_RE.match(head) else normalise_name(
            path.parent.name.replace("_", " "))
        ref = stored.get((key, path.name))
        if ref is None or ref.shape != native[i].shape:
            continue
        ref = ref / max(float(np.linalg.norm(ref)), 1e-6)
        sims.append(float(native[i] @ ref))
        paired_mine.append(native[i])
        paired_ref.append(ref)
        paired_actor.append(actor_of[i])
    if not sims:
        print(f"[verify] no overlap with {gallery_path} — nothing checked",
              file=sys.stderr)
        return

    sims_arr = np.asarray(sims)
    print(f"[verify] {len(sims)} embeddings vs {gallery_path.name}: "
          f"mean cos={sims_arr.mean():.4f}  min={sims_arr.min():.4f}",
          file=sys.stderr)
    act = np.asarray(paired_actor)
    for label, mat in (("this script", np.asarray(paired_mine)),
                       ("stored gallery", np.asarray(paired_ref))):
        pos, neg, d = _separation(mat, act)
        print(f"[verify]   {label:>14s}: same-actor {pos:.3f}  "
              f"different-actor {neg:.3f}  d'={d:.2f}", file=sys.stderr)
    if sims_arr.mean() < 0.99:
        print("[verify]   embeddings differ from the stored gallery. If d' is "
              "comparable this is a backend difference (e.g. a TensorRT fp16 "
              "build), not a broken port; the study is self-consistent either "
              "way. If d' collapsed, the port is wrong.", file=sys.stderr)


def err(msg: str) -> int:
    print(f"error: {msg}", file=sys.stderr)
    return 2


if __name__ == "__main__":
    sys.exit(main())
