#!/usr/bin/env python3
"""Landmark voting: average SCRFD's overlapping detections instead of discarding them.

SCRFD predicts a face from many anchors; NMS keeps the single highest-scoring
box and throws the rest away. Each discarded box carries its own 5-landmark
estimate of the SAME face, so the survivors are one sample from a distribution
we could be averaging over.

    baseline  conf 0.50, nms 0.40 — the shipped settings, one box per face
    voted     conf 0.30, nms 0.90 — duplicates survive, then grouped by IoU and
              the 5 landmarks averaged, weighted by detection confidence

Why this is worth trying when the mesh failed: the mesh moved the landmarks off
the definition ArcFace was trained on (a lip-ring centroid is not an annotated
mouth corner), and the embedder punished it. A confidence-weighted mean of
SCRFD's OWN landmark predictions is the same kind of point, just with less
variance — it should stay on-distribution while being steadier.

Scored on cross-clip identification through the production sigmoid, which is
the thing that actually broke. Raw similarity shown only to locate the
threshold; it decides nothing.

    LD_PRELOAD=/usr/lib/libcudnn_cnn.so.9 python3 landmark_voting.py
"""
import sys, glob, json, os

sys.path.insert(0, "/home/dtourolle/Development/Jray-project/scene-actor-extraction/build-ort")
import sae_embed          # before cv2 — see alignment_compare.py

import numpy as np
import cv2

import argparse

ROOT = "/home/dtourolle/Development/Jray-project/scene-actor-extraction/"
M = ROOT + "models/"
CLIPS = ["5157344", "5157339"]
PROB_THRESHOLD = 0.754
GROUP_IOU = 0.55          # detections overlapping this much are the same face
MATCH_IOU = 0.35          # tie a detection to the hand-labelled face

_ap = argparse.ArgumentParser()
_ap.add_argument("--detector", default="scrfd_500m_bnkps.onnx",
                 help="detector under models/. SCRFD sizes 500m / 2.5g / 10g come "
                      "from InsightFace's buffalo_sc / buffalo_m / buffalo_l packs")
_ap.add_argument("--vote-conf", type=float, default=None,
                 help="confidence floor for the voting pass. Omit to auto-tune "
                      "it to --target-votes")
_ap.add_argument("--target-votes", type=int, default=3,
                 help="votes per face to tune --vote-conf towards, so detectors "
                      "are compared at equal redundancy rather than equal settings")
_args = _ap.parse_args()

base_eng = sae_embed.FaceEmbedder(detector_model=M + _args.detector,
                                  arcface_model=M + "LVFace-B_Glint360K.onnx",
                                  conf=0.5, nms=0.4, max_side=0)


def _make_vote_engine(conf):
    # Same models, looser suppression: keep the duplicates NMS would have removed.
    return sae_embed.FaceEmbedder(detector_model=M + _args.detector,
                                  arcface_model=M + "LVFace-B_Glint360K.onnx",
                                  conf=conf, nms=0.9, max_side=0)


cal = sae_embed.gallery_calibration(ROOT + "gallery_lvface.h5")
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
    i = (x1 - x0) * (y1 - y0)
    return i / (aw * ah + bw * bh - i)


def vote(dets):
    """Group overlapping detections, return (bbox, landmarks, conf, n_votes)."""
    items = sorted(dets, key=lambda d: -d.confidence)
    used, out = [False] * len(items), []
    for i, d in enumerate(items):
        if used[i]:
            continue
        grp = [d]
        used[i] = True
        for j in range(i + 1, len(items)):
            if not used[j] and iou(list(d.bbox), list(items[j].bbox)) >= GROUP_IOU:
                used[j] = True
                grp.append(items[j])
        w = np.array([g.confidence for g in grp], dtype=np.float32)
        w = w / w.sum()
        lms = np.stack([np.array(g.landmarks, dtype=np.float32).reshape(5, 2) for g in grp])
        bxs = np.stack([np.array(list(g.bbox), dtype=np.float32) for g in grp])
        out.append((( w[:, None] * bxs).sum(0), (w[:, None, None] * lms).sum(0),
                    float(grp[0].confidence), len(grp)))
    return out


def tune_vote_conf(target, sample=6):
    """Pick the confidence floor giving ~target detections per face to average.

    A larger SCRFD is more confident and suppresses harder, so at a fixed floor
    it emits fewer overlapping anchors — median 2 against 500m's 3. Comparing
    detectors at equal SETTINGS therefore also compares them at unequal
    redundancy, and the voting arm is handicapped for the bigger models. Tuning
    each to the same votes-per-face isolates landmark quality from how much
    there was to average.
    """
    frames = sorted(glob.glob(f"frames/d{CLIPS[0]}_*.png"))[:sample]
    imgs = [cv2.imread(f) for f in frames]
    best = (None, None, 1e9)
    for conf in (0.30, 0.20, 0.12, 0.07, 0.04, 0.02, 0.01):
        eng = _make_vote_engine(conf)
        sizes = [n for img in imgs for _, _, _, n in vote(eng.detect(img))]
        if not sizes:
            continue
        med = float(np.median(sizes))
        if abs(med - target) < best[2]:
            best = (conf, eng, abs(med - target))
        print(f"[tune] conf={conf:.2f} -> median {med:.0f} votes/face", file=sys.stderr)
        if med >= target:
            break
    if best[1] is None:
        print(f"[tune] no confidence floor reached {target} votes/face; "
              f"falling back to 0.30", file=sys.stderr)
        return 0.30, _make_vote_engine(0.30)
    print(f"[tune] chose conf={best[0]:.2f} for ~{target} votes/face", file=sys.stderr)
    return best[0], best[1]


if _args.vote_conf is not None:
    VOTE_CONF, vote_eng = _args.vote_conf, _make_vote_engine(_args.vote_conf)
else:
    VOTE_CONF, vote_eng = tune_vote_conf(_args.target_votes)

cal = sae_embed.gallery_calibration(ROOT + "gallery_lvface.h5")


def collect(clip):
    lab = {os.path.basename(p): os.path.basename(os.path.dirname(p))
           for p in glob.glob(f"labelling/{clip}/*/*.jpg")
           if os.path.basename(os.path.dirname(p)) not in ("discard", "unsorted")}
    man = {m["file"]: m for m in json.load(open(f"labelling/{clip}/manifest.json"))}
    rows, votes = [], []
    for frame in sorted({man[f]["frame"] for f in lab}):
        img = cv2.imread(f"frames/d{clip}_{frame}.png")
        base = base_eng.detect(img)
        voted = vote(vote_eng.detect(img))
        for fname, person in lab.items():
            m = man[fname]
            if m["frame"] != frame or m["idx"] >= len(base):
                continue
            d = base[m["idx"]]
            lm5 = np.array(d.landmarks, dtype=np.float32).reshape(5, 2)
            c_b = sae_embed.align_face(img, lm5)

            # the voted group covering the same face
            best, best_v = None, 0.0
            for bbox, lms, conf, n in voted:
                v = iou(list(bbox), list(d.bbox))
                if v > best_v:
                    best_v, best = v, (lms, n)
            c_v = None
            if best and best_v >= MATCH_IOU:
                c_v = sae_embed.align_face(img, best[0].astype(np.float32))
                votes.append(best[1])
            rec = {"person": person}
            rec["base"] = np.asarray(base_eng.embed_crop(c_b), np.float32) if c_b is not None else None
            rec["voted"] = np.asarray(base_eng.embed_crop(c_v), np.float32) if c_v is not None else None
            rows.append(rec)
    return rows, votes


data, allv = {}, []
for c in CLIPS:
    data[c], v = collect(c)
    allv += v
    print(f"[{c}] {len(data[c])} crops", file=sys.stderr)
print(f"[voting] group size: median {np.median(allv):.0f}, "
      f"mean {np.mean(allv):.1f}, max {max(allv)} detections averaged per face",
      file=sys.stderr)

GAL, PRB = "5157344", "5157339"
print(f"\ndetector={_args.detector}  vote_conf={VOTE_CONF:.2f}  "
      f"gallery {GAL} -> probe {PRB}, P>{PROB_THRESHOLD}\n")
print(f"{'align':>8}{'person':>8}{'n_gal':>7}{'n_prb':>7}"
      f"{'within-clip':>13}{'cross-clip':>12}{'hit rate':>10}")
summary = {}
for key in ("base", "voted"):
    gal, prb = {}, {}
    for r in data[GAL]:
        if r[key] is not None:
            gal.setdefault(r["person"], []).append(r[key])
    for r in data[PRB]:
        if r[key] is not None:
            prb.setdefault(r["person"], []).append(r[key])
    gal = {p: np.stack(v) for p, v in gal.items()}
    prb = {p: np.stack(v) for p, v in prb.items()}
    hits = tot = 0
    for p in sorted(set(gal) & set(prb)):
        pp = prb[p] @ prb[p].T
        np.fill_diagonal(pp, -1)
        within = float(np.median(pp.max(axis=1))) if len(pp) > 1 else float("nan")
        cross = float(np.median((gal[p] @ prb[p].T).max(axis=0)))
        h = 0
        for e in prb[p]:
            bp, bn = 0.0, None
            for q in gal:
                v = cal.probability(float((gal[q] @ e).max()))
                if v > bp:
                    bp, bn = v, q
            if bp > PROB_THRESHOLD and bn == p:
                h += 1
        hits += h; tot += len(prb[p])
        print(f"{key:>8}{p:>8}{len(gal[p]):>7}{len(prb[p]):>7}"
              f"{cal.probability(within):>6.3f}/{within:<6.3f}"
              f"{cal.probability(cross):>6.3f}/{cross:<5.3f}{100*h/len(prb[p]):>9.0f}%")
    summary[key] = (hits, tot)
    print(f"{key:>8}{'ALL':>8}{'':>14}{'':>25}{100*hits/max(tot,1):>9.0f}%\n")

hb, tb = summary["base"]; hv, tv = summary["voted"]
print(f"voting vs baseline: {100*hv/max(tv,1) - 100*hb/max(tb,1):+.1f} points "
      f"of cross-clip TPI ({hb}/{tb} -> {hv}/{tv})")