#!/usr/bin/env python3
"""Impact of input resolution on cross-source identification.

Gallery is built from one clip at NATIVE resolution. Probes come from the other
clip with the WHOLE FRAME downscaled before it reaches the detector, so
detection and landmark regression degrade together with the pixels. That is the
measurement VR-005 structurally could not make: it degraded an already-aligned
112x112 crop, holding alignment perfect, so it isolated the embedder's
resolution sensitivity and excluded everything upstream of it.

    python3 resolution_sweep.py [--gallery-clip 5157339] [--detector scrfd_500m_bnkps.onnx]

Ground truth
------------
Hand-sorted person folders. Probe detections at reduced scale are tied back to
a labelled face GEOMETRICALLY — the box is mapped to native coordinates and
matched by IoU. Never by embedding similarity, which would be circular: it
would keep the faces the embedder still gets right and silently drop the ones
this sweep exists to find.

A probe whose label is only in the probe clip is OUT OF GALLERY. Naming it is a
true out-of-cast misID, the error the per-scene scorer weights 10x, so it is
counted separately from naming the wrong gallery member.

Metric
------
The calibrated probability from the PRODUCTION gallery sigmoid, never a raw
cosine (AR-024). Per-actor best-of-N similarity -> probability -> accept above
prob_threshold. This is identification, so the matcher's prior applies;
config.hpp has match_prior 0.5, i.e. log_prior_odds = 0.

Everything runs through the shipped C++ via sae_embed.
"""
import sys, glob, json, os, argparse
import numpy as np
import cv2

sys.path.insert(0, "/home/dtourolle/Development/Jray-project/scene-actor-extraction/build-ort")
import sae_embed

ROOT = "/home/dtourolle/Development/Jray-project/scene-actor-extraction/"
M = ROOT + "models/"
PROB_THRESHOLD = 0.754      # config.hpp:67
LOG_PRIOR_ODDS = 0.0        # config.hpp:61 match_prior=0.5
IOU_MIN = 0.3               # geometric label carry-down
SCALES = [1.0, 0.8, 0.6, 0.5, 0.4, 0.3, 0.25, 0.2, 0.15, 0.12, 0.09, 0.06]

ap = argparse.ArgumentParser()
ap.add_argument("--gallery-clip", default="5157339")
ap.add_argument("--probe-clip", default="5157344")
ap.add_argument("--detector", default="scrfd_500m_bnkps.onnx")
ap.add_argument("--embedder", default="LVFace-B_Glint360K.onnx")
ap.add_argument("--gallery-calibration", default=ROOT + "gallery_lvface.h5")
ap.add_argument("--out", default="results_resolution_sweep.json")
args = ap.parse_args()

eng = sae_embed.FaceEmbedder(detector_model=M + args.detector,
                             arcface_model=M + args.embedder,
                             conf=0.5, nms=0.4, max_side=0)
cal = sae_embed.gallery_calibration(args.gallery_calibration)
print(f"[calibration] global: {cal}", file=sys.stderr)


def labelled(clip):
    """{filename: person} from the hand-sorted folders, ignoring discard."""
    out = {}
    for path in glob.glob(f"labelling/{clip}/*/*.jpg"):
        person = os.path.basename(os.path.dirname(path))
        if person in ("discard", "unsorted"):
            continue
        out[os.path.basename(path)] = person
    return out


def manifest(clip):
    return {m["file"]: m for m in json.load(open(f"labelling/{clip}/manifest.json"))}


def iou(a, b):
    ax, ay, aw, ah = a; bx, by, bw, bh = b
    x0, y0 = max(ax, bx), max(ay, by)
    x1, y1 = min(ax + aw, bx + bw), min(ay + ah, by + bh)
    if x1 <= x0 or y1 <= y0:
        return 0.0
    inter = (x1 - x0) * (y1 - y0)
    return inter / (aw * ah + bw * bh - inter)


# ── gallery: native resolution, labelled faces only ──────────────────────────
g_lab, g_man = labelled(args.gallery_clip), manifest(args.gallery_clip)
gal = {}
for frame in sorted({g_man[f]["frame"] for f in g_lab}):
    img = cv2.imread(f"frames/d{args.gallery_clip}_{frame}.png")
    dets = eng.detect(img)
    for fname, person in g_lab.items():
        m = g_man[fname]
        if m["frame"] != frame or m["idx"] >= len(dets):
            continue
        lm = np.array(dets[m["idx"]].landmarks, dtype=np.float32).reshape(5, 2)
        crop = sae_embed.align_face(img, lm)
        if crop is None:
            continue
        gal.setdefault(person, []).append(np.asarray(eng.embed_crop(crop), dtype=np.float32))
gal = {p: np.stack(v) for p, v in gal.items() if v}
people = sorted(gal)
print(f"[gallery] {args.gallery_clip} @native: "
      f"{ {p: len(v) for p, v in gal.items()} }", file=sys.stderr)

# ── probe ground truth at native resolution ──────────────────────────────────
p_lab, p_man = labelled(args.probe_clip), manifest(args.probe_clip)
truth = {}          # frame -> [(bbox_native, person)]
for fname, person in p_lab.items():
    m = p_man[fname]
    truth.setdefault(m["frame"], []).append((m["bbox"], person))
n_out = sum(1 for p in set(p_lab.values()) if p not in people)
print(f"[probe] {args.probe_clip}: {len(p_lab)} labelled faces, "
      f"{len(set(p_lab.values()))} people, {n_out} of them out-of-gallery",
      file=sys.stderr)

# ── sweep ────────────────────────────────────────────────────────────────────
print(f"\n{'scale':>6}{'frame':>11}{'face px':>9}{'found':>7}{'matched':>9}"
      f"{'TPI':>8}{'FPI-in':>8}{'FPI-out':>9}{'TBI':>8}")
results = []
for s in SCALES:
    tpi = fpi_in = fpi_out = tbi = 0
    n_found = n_matched = 0
    pxs = []
    for frame, gts in sorted(truth.items()):
        img = cv2.imread(f"frames/d{args.probe_clip}_{frame}.png")
        if s != 1.0:
            img = cv2.resize(img, None, fx=s, fy=s, interpolation=cv2.INTER_AREA)
        dets = eng.detect(img)
        n_found += len(dets)
        for d in dets:
            x, y, w, h = d.bbox
            native = (x / s, y / s, w / s, h / s)       # geometric carry-down
            best, best_iou = None, 0.0
            for gt_box, person in gts:
                v = iou(native, gt_box)
                if v > best_iou:
                    best_iou, best = v, person
            if best_iou < IOU_MIN:
                continue                                 # spurious / unlabelled
            n_matched += 1
            pxs.append(min(w, h))

            lm = np.array(d.landmarks, dtype=np.float32).reshape(5, 2)
            crop = sae_embed.align_face(img, lm)
            if crop is None:
                tbi += 1                                 # degenerate alignment
                continue
            emb = np.asarray(eng.embed_crop(crop), dtype=np.float32)
            best_p, best_name = 0.0, None
            for p in people:                             # per-actor best-of-N
                prob = cal.probability(float((gal[p] @ emb).max()), LOG_PRIOR_ODDS)
                if prob > best_p:
                    best_p, best_name = prob, p
            if best_p <= PROB_THRESHOLD:
                tbi += 1
            elif best not in people:
                fpi_out += 1                             # named someone absent from the gallery
            elif best_name == best:
                tpi += 1
            else:
                fpi_in += 1
    n = max(1, n_matched)
    med_px = float(np.median(pxs)) if pxs else 0.0
    print(f"{s:>6.2f}{f'{int(4096*s)}x{int(2160*s)}':>11}{med_px:>9.0f}"
          f"{n_found:>7}{n_matched:>9}"
          f"{100*tpi/n:>7.1f}%{100*fpi_in/n:>7.1f}%{100*fpi_out/n:>8.1f}%{100*tbi/n:>7.1f}%")
    results.append({"scale": s, "median_face_px": med_px, "detections": n_found,
                    "matched_to_truth": n_matched, "tpi_pct": 100*tpi/n,
                    "fpi_in_gallery_pct": 100*fpi_in/n, "fpi_out_of_gallery_pct": 100*fpi_out/n,
                    "tbi_pct": 100*tbi/n})

json.dump({"gallery_clip": args.gallery_clip, "probe_clip": args.probe_clip,
           "detector": args.detector, "embedder": args.embedder,
           "prob_threshold": PROB_THRESHOLD, "log_prior_odds": LOG_PRIOR_ODDS,
           "calibration": {"a": cal.a, "b": cal.b},
           "gallery_people": people, "results": results},
          open(args.out, "w"), indent=2)
print(f"\nwrote {args.out}", file=sys.stderr)