#!/usr/bin/env python3
"""Dump face crops from both clips for hand-labelling.

Writes labelling/<clip>/unsorted/<name>.jpg — a context crop around each
detection, big enough to recognise a person by eye. Move them into
labelling/<clip>/person_A/, person_B/, ... and the sweep reads those folders as
ground truth.

Filenames carry a cNN_ cluster-hint prefix so visually similar faces sort next
to each other in a file manager. The hint is only an ordering convenience —
the folder you drop a file into is what counts, and the sweep never reads the
prefix.

Detection and alignment run through the shipped C++ (sae_embed). Every crop
keeps its clip, frame and native-resolution bbox in manifest.json, so probe
detections at reduced scale can be tied back to a labelled face geometrically,
by position, rather than by embedding similarity — which would be circular.
"""
import sys, glob, json, os, shutil
import numpy as np
import cv2

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

M = "/home/dtourolle/Development/Jray-project/scene-actor-extraction/models/"
CLIPS = ["5157339", "5157344"]
MIN_PX = 60
CTX = 256          # context-crop side, for human recognisability

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

for clip in CLIPS:
    out_dir = f"labelling/{clip}/unsorted"
    if os.path.isdir(f"labelling/{clip}"):
        print(f"[skip] labelling/{clip} exists — not overwriting your sorting",
              file=sys.stderr)
        continue
    os.makedirs(out_dir, exist_ok=True)

    entries = []
    for p in sorted(glob.glob(f"pex/d{clip}_*.png")):
        frame = p.rsplit("_", 1)[-1].split(".")[0]
        img = cv2.imread(p)
        for i, d in enumerate(eng.detect(img)):
            x, y, w, h = d.bbox
            if min(w, h) < MIN_PX:
                continue
            lm = np.array(d.landmarks, dtype=np.float32).reshape(5, 2)
            crop = sae_embed.align_face(img, lm)
            if crop is None:
                continue
            emb = np.asarray(eng.embed_crop(crop), dtype=np.float32)

            pad = int(0.5 * max(w, h))
            x0, y0 = max(0, int(x) - pad), max(0, int(y) - pad)
            x1, y1 = min(img.shape[1], int(x + w) + pad), min(img.shape[0], int(y + h) + pad)
            ctx = cv2.resize(img[y0:y1, x0:x1], (CTX, CTX))

            entries.append({"clip": clip, "frame": frame, "idx": i,
                            "bbox": [float(x), float(y), float(w), float(h)],
                            "px": float(min(w, h)), "conf": float(d.confidence),
                            "emb": emb, "ctx": ctx})

    # cluster hint only — greedy, purely to group similar faces in the file list
    E = np.stack([e["emb"] for e in entries])
    hint = -np.ones(len(entries), int)
    k = 0
    for i in range(len(entries)):
        if hint[i] >= 0:
            continue
        hint[i] = k
        for j in range(i + 1, len(entries)):
            if hint[j] < 0 and float(E[i] @ E[j]) > 0.5:
                hint[j] = k
        k += 1

    manifest = []
    for e, h in zip(entries, hint):
        name = f"c{h:02d}_{e['clip']}_f{e['frame']}_i{e['idx']}_{int(e['px'])}px.jpg"
        cv2.imwrite(f"{out_dir}/{name}", e["ctx"])
        manifest.append({k: v for k, v in e.items() if k not in ("emb", "ctx")}
                        | {"file": name, "cluster_hint": int(h)})

    json.dump(manifest, open(f"labelling/{clip}/manifest.json", "w"), indent=1)
    print(f"[{clip}] {len(manifest)} crops in {out_dir}, {k} cluster hints, "
          f"face px {min(m['px'] for m in manifest):.0f}–{max(m['px'] for m in manifest):.0f}",
          file=sys.stderr)