#!/usr/bin/env python3
"""Integrity check on the labelled set, before it is used as ground truth.

TRACES: VR-013 | PR-002

VR-013's ground truth is hand-sorted rather than propagated by embedding
similarity, because propagation would keep only the faces the embedder already
gets right and silently drop the ones the sweep exists to find. This script is
what makes that claim checkable, so it is part of the requirement rather than a
helper of it.

Checks, loudest failure first:

1. INDEX INTEGRITY. Each crop's embedding is taken by re-detecting its source
   frame and indexing with the manifest's `idx`. If detection order is not
   reproducible, the thumbnail you sorted and the embedding that gets scored
   are different faces — you would see a correct picture and score the wrong
   person, with nothing to signal it. Every crop's re-detected bbox is compared
   against the manifest's.

2. NO CROP IN TWO FOLDERS, and every manifest entry accounted for — so a
   move that half-completed cannot silently duplicate or drop a label.

3. ALIGNMENT. The 112x112 warp is what the embedder actually sees; the
   thumbnail is only context for your eyes. verify_<person>.jpg pairs them:
   context-with-box on top, the real aligned crop beneath. A profile face whose
   alignment has collapsed is obvious there and nowhere else.

4. SEPARATION. Per person, the calibrated P of their own crops against the
   other people's, using the global gallery sigmoid. A label set where someone
   matches another person better than themselves is mislabelled.

Nothing here changes a label. It reports.
"""
import sys, glob, json, os
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/"
EMBEDDER = M + "LVFace-B_Glint360K.onnx"
GALLERY  = ROOT + "gallery_lvface.h5"
CLIPS = ["5157344", "5157339"]
THUMB = 130
COLS = 10

eng = sae_embed.FaceEmbedder(detector_model=M + "scrfd_500m_bnkps.onnx",
                             arcface_model=EMBEDDER, conf=0.5, nms=0.4, max_side=0)

fail = 0
rows = []

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

    # where each crop currently sits -> its label
    placed = {}
    for path in glob.glob(f"labelling/{clip}/*/*.jpg"):
        person = os.path.basename(os.path.dirname(path))
        if person in ("discard", "unsorted"):
            continue   # not people; scoring them would invent an extra identity
        fname = os.path.basename(path)
        if fname in placed:
            print(f"[FAIL] {fname} appears in both {placed[fname][0]} and {person}")
            fail += 1
        placed[fname] = (person, path)

    missing = set(man) - set(placed)
    extra = set(placed) - set(man)
    if missing:
        print(f"[warn] {clip}: {len(missing)} manifest crops not in any folder")
    if extra:
        print(f"[FAIL] {clip}: {len(extra)} files with no manifest entry: "
              f"{sorted(extra)[:3]}")
        fail += 1

    # index integrity + alignment, frame by frame
    by_frame = {}
    for fname, (person, path) in placed.items():
        if fname in man:
            by_frame.setdefault(man[fname]["frame"], []).append((fname, person, path))

    bad_idx = 0
    for frame, items in sorted(by_frame.items()):
        img = cv2.imread(f"frames/d{clip}_{frame}.png")
        if img is None:
            print(f"[FAIL] missing frames/d{clip}_{frame}.png")
            fail += 1
            continue
        dets = eng.detect(img)
        for fname, person, path in items:
            m = man[fname]
            i = m["idx"]
            if i >= len(dets):
                print(f"[FAIL] {fname}: idx {i} >= {len(dets)} detections now")
                bad_idx += 1
                continue
            got = [float(v) for v in dets[i].bbox]
            want = m["bbox"]
            if max(abs(a - b) for a, b in zip(got, want)) > 1.0:
                print(f"[FAIL] {fname}: manifest bbox {[round(v) for v in want]} "
                      f"!= re-detected {[round(v) for v in got]}")
                bad_idx += 1
                continue
            lm = np.array(dets[i].landmarks, dtype=np.float32).reshape(5, 2)
            crop = sae_embed.align_face(img, lm)
            if crop is None:
                print(f"[warn] {fname}: alignment degenerate, no crop reaches the embedder")
                continue
            rows.append({"clip": clip, "person": person, "file": fname, "path": path,
                         "px": m["px"], "aligned": np.asarray(crop),
                         "emb": np.asarray(eng.embed_crop(crop), dtype=np.float32)})
    fail += bad_idx
    print(f"[{clip}] {len(placed)} placed, {len(by_frame)} frames, "
          f"index mismatches: {bad_idx}")

if not rows:
    sys.exit("nothing to verify")

# ── separation, through the global gallery sigmoid ───────────────────────────
cal = sae_embed.gallery_calibration(GALLERY)
E = np.stack([r["emb"] for r in rows])
people = sorted({r["person"] for r in rows})
lab = np.array([people.index(r["person"]) for r in rows])
S = E @ E.T
np.fill_diagonal(S, -1.0)

print(f"\n{'person':>8}{'crops':>7}{'344':>6}{'339':>6}"
      f"{'P(self)':>10}{'P(other)':>10}{'worst':>8}")
for k, p in enumerate(people):
    mine = np.where(lab == k)[0]
    if len(mine) < 2:
        continue
    self_sim = S[np.ix_(mine, mine)].max(axis=1)
    other_sim = S[np.ix_(mine, np.where(lab != k)[0])].max(axis=1)
    p_self = np.array([cal.probability(float(s)) for s in self_sim])
    p_other = np.array([cal.probability(float(s)) for s in other_sim])
    n344 = sum(1 for i in mine if rows[i]["clip"] == "5157344")
    n339 = len(mine) - n344
    # a crop that matches someone else better than anyone of its own label
    worst = int((other_sim > self_sim).sum())
    print(f"{p:>8}{len(mine):>7}{n344:>6}{n339:>6}"
          f"{np.median(p_self):>10.3f}{np.median(p_other):>10.3f}{worst:>8}")
    if worst:
        for i in mine[other_sim > self_sim]:
            print(f"          suspect: {rows[i]['file']}  "
                  f"P(self)={cal.probability(float(self_sim[list(mine).index(i)])):.3f} "
                  f"< P(other)={cal.probability(float(other_sim[list(mine).index(i)])):.3f}")

# ── verify sheets: context+box over the actual aligned crop ──────────────────
for p in people:
    items = [r for r in rows if r["person"] == p]
    items.sort(key=lambda r: (r["clip"], r["file"]))
    n = len(items)
    sheet_rows = (n + COLS - 1) // COLS
    H = THUMB * 2 + 22
    sheet = np.full((sheet_rows * H, COLS * THUMB, 3), 25, np.uint8)
    for j, r in enumerate(items):
        rr, cc = divmod(j, COLS)
        y, x = rr * H, cc * THUMB
        ctx = cv2.imread(r["path"])
        if ctx is not None:
            sheet[y:y + THUMB, x:x + THUMB] = cv2.resize(ctx, (THUMB, THUMB))
        sheet[y + THUMB:y + 2 * THUMB, x:x + THUMB] = cv2.resize(r["aligned"], (THUMB, THUMB))
        cv2.putText(sheet, f"{r['clip'][-3:]} {int(r['px'])}px",
                    (x + 3, y + 2 * THUMB + 15),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.38, (150, 220, 150), 1)
    cv2.imwrite(f"labelling/verify_{p}.jpg", sheet)
    print(f"  verify_{p}.jpg: {n} crops (top row context, bottom row what the embedder sees)")

print(f"\n{'PASS' if fail == 0 else f'{fail} FAILURES'}")
sys.exit(1 if fail else 0)