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

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)