#!/usr/bin/env python3
"""What explains the misses? Head pose, face size, blur, detector confidence.

For every hand-labelled probe face, computes the calibrated probability against
its OWN gallery entry — so a low value is a false negative, not a mistake about
who it is — and pairs it with covariates that might explain the failure.

Head pose comes from solvePnP of the 5 landmarks against a canonical 3D face,
giving yaw/pitch/roll in degrees.

    CAVEAT, and it matters: the pose estimate is derived from the same 5
    landmarks the alignment uses. Where those landmarks are unreliable the pose
    estimate is unreliable too, and both degrade for the same reason. So this
    can show that failures concentrate at high yaw; it cannot cleanly separate
    "the head was turned" from "the landmarks were wrong because the head was
    turned". Those are the same physical cause, but not the same fix — the
    first argues for gallery pose coverage, the second for a better landmark
    source.

A sanity check is printed first: pose is estimated per person, and if it does
not recover what is visible in the review sheets (one subject frontal, another
in profile, another looking down) then the estimate is not worth reading.

Similarities go through the production gallery sigmoid, never compared raw.
"""
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/"
GALLERY_CLIP, PROBE_CLIP = "5157344", "5157339"
PROB_THRESHOLD = 0.754

# Canonical 3D face, ordered as types.hpp:60 —
# [0] right-eye [1] left-eye [2] nose [3] right-mouth [4] left-mouth.
# The subject's right eye sits to the LEFT in image space, hence the negative X.
FACE_3D = np.array([
    (-34.0,  35.0, -28.0),
    ( 34.0,  35.0, -28.0),
    (  0.0,   0.0,   0.0),
    (-26.0, -32.0, -25.0),
    ( 26.0, -32.0, -25.0),
], dtype=np.float64)

eng = sae_embed.FaceEmbedder(detector_model=M + "scrfd_500m_bnkps.onnx",
                             arcface_model=M + "LVFace-B_Glint360K.onnx",
                             conf=0.5, nms=0.4, max_side=0)
cal = sae_embed.gallery_calibration(ROOT + "gallery_lvface.h5")


def head_pose(lm, w, h):
    """yaw, pitch, roll in degrees. Focal length assumed = image width."""
    cam = np.array([[w, 0, w / 2], [0, w, h / 2], [0, 0, 1]], dtype=np.float64)
    ok, rvec, _ = cv2.solvePnP(FACE_3D, lm.astype(np.float64), cam, None,
                               flags=cv2.SOLVEPNP_EPNP)
    if not ok:
        return None
    R, _ = cv2.Rodrigues(rvec)
    sy = np.sqrt(R[0, 0] ** 2 + R[1, 0] ** 2)
    if sy > 1e-6:
        pitch = np.degrees(np.arctan2(-R[2, 0], sy))
        yaw = np.degrees(np.arctan2(R[1, 0], R[0, 0]))
        roll = np.degrees(np.arctan2(R[2, 1], R[2, 2]))
    else:
        pitch = np.degrees(np.arctan2(-R[2, 0], sy)); yaw = 0.0
        roll = np.degrees(np.arctan2(-R[1, 2], R[1, 1]))
    # solvePnP's yaw wraps near +/-180 for a face pointing at the camera;
    # fold it to a "degrees away from frontal" magnitude.
    yaw = ((yaw + 180) % 360) - 180
    if abs(yaw) > 90:
        yaw = np.sign(yaw) * (180 - abs(yaw))
    return yaw, pitch, roll


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 = []
    for frame in sorted({man[f]["frame"] for f in lab}):
        img = cv2.imread(f"frames/d{clip}_{frame}.png")
        dets = eng.detect(img)
        H, W = img.shape[:2]
        for f, person in lab.items():
            m = man[f]
            if m["frame"] != frame or m["idx"] >= len(dets):
                continue
            d = dets[m["idx"]]
            lm = np.array(d.landmarks, dtype=np.float32).reshape(5, 2)
            crop = sae_embed.align_face(img, lm)
            if crop is None:
                continue
            pose = head_pose(lm, W, H)
            x, y, w, h = d.bbox
            g = cv2.cvtColor(np.asarray(crop), cv2.COLOR_BGR2GRAY)
            rows.append({
                "person": person, "px": float(min(w, h)), "conf": float(d.confidence),
                "yaw": pose[0] if pose else np.nan, "pitch": pose[1] if pose else np.nan,
                "roll": pose[2] if pose else np.nan,
                "blur": float(cv2.Laplacian(g, cv2.CV_64F).var()),
                "emb": np.asarray(eng.embed_crop(crop), dtype=np.float32)})
    return rows


gal_rows = collect(GALLERY_CLIP)
prb_rows = collect(PROBE_CLIP)
gal = {}
for r in gal_rows:
    gal.setdefault(r["person"], []).append(r["emb"])
gal = {p: np.stack(v) for p, v in gal.items()}

for r in prb_rows:
    if r["person"] in gal:
        s = float((gal[r["person"]] @ r["emb"]).max())     # best-of-N, own actor
        r["p"] = cal.probability(s)
        r["sim"] = s
    else:
        r["p"] = np.nan
rows = [r for r in prb_rows if not np.isnan(r.get("p", np.nan))]
print(f"[data] {len(rows)} labelled probe faces with a gallery entry\n", file=sys.stderr)

# ── sanity check: does the pose estimate recover what the sheets show? ───────
print("pose by person (does this match the review sheets?)")
print(f"{'person':>7}{'n':>5}{'|yaw| med':>11}{'pitch med':>11}{'P med':>8}{'hit rate':>10}")
for p in sorted({r['person'] for r in rows}):
    sub = [r for r in rows if r["person"] == p]
    print(f"{p:>7}{len(sub):>5}"
          f"{np.median([abs(r['yaw']) for r in sub]):>11.1f}"
          f"{np.median([r['pitch'] for r in sub]):>11.1f}"
          f"{np.median([r['p'] for r in sub]):>8.3f}"
          f"{100*np.mean([r['p'] > PROB_THRESHOLD for r in sub]):>9.0f}%")

# ── P binned by each covariate ───────────────────────────────────────────────
def binned(name, key, edges, fmt="{:.0f}"):
    print(f"\nP(match) by {name}")
    print(f"{'bin':>16}{'n':>5}{'P med':>9}{'hit rate':>10}{'sim med':>9}")
    vals = np.array([r[key] for r in rows])
    for lo, hi in zip(edges[:-1], edges[1:]):
        sub = [r for r, v in zip(rows, vals) if lo <= v < hi]
        if not sub:
            continue
        lbl = f"{fmt.format(lo)}–{fmt.format(hi)}"
        print(f"{lbl:>16}{len(sub):>5}"
              f"{np.median([r['p'] for r in sub]):>9.3f}"
              f"{100*np.mean([r['p'] > PROB_THRESHOLD for r in sub]):>9.0f}%"
              f"{np.median([r['sim'] for r in sub]):>9.3f}")

for r in rows:
    r["absyaw"] = abs(r["yaw"])
    r["abspitch"] = abs(r["pitch"])
binned("|yaw| (deg from frontal)", "absyaw", [0, 10, 20, 30, 45, 60, 91])
binned("|pitch| (deg)", "abspitch", [0, 10, 20, 30, 45, 91])
binned("face size (px)", "px", [0, 130, 150, 175, 200, 400])
binned("blur (laplacian var)", "blur", [0, 50, 150, 400, 1000, 1e9])
binned("detector confidence", "conf", [0.5, 0.6, 0.7, 0.8, 0.9, 1.01], "{:.2f}")

# ── how much does each covariate actually explain? ───────────────────────────
print("\nSpearman rank correlation with P(match):")
def spearman(a, b):
    ra = np.argsort(np.argsort(a)); rb = np.argsort(np.argsort(b))
    return float(np.corrcoef(ra, rb)[0, 1])
P = np.array([r["p"] for r in rows])
for key, label in [("absyaw", "|yaw|"), ("abspitch", "|pitch|"), ("px", "face px"),
                   ("blur", "blur"), ("conf", "detector conf")]:
    v = np.array([r[key] for r in rows])
    print(f"   {label:>14}: {spearman(v, P):+.3f}")

json.dump([{k: v for k, v in r.items() if k != "emb"} for r in rows],
          open("failure_analysis.json", "w"), indent=1, default=float)