#!/usr/bin/env python3
"""Estimate head pose per crop, and build a page to confirm or correct it.

Why not solvePnP on the 5 detector landmarks: those landmarks collapse on
turned faces, so the estimator breaks precisely on the crops whose pose we care
about. Run that way it reported the profile subject as the MOST frontal of the
four, which is how we know not to trust it.

Instead the estimate comes from the MediaPipe face mesh (468 points, run via
OpenCV DNN — the same model rPPG-kahn uses) and a symmetry measure that needs
no 3D model:

    yaw_ratio = (dL - dR) / (dL + dR)

over left/right symmetric vertex pairs, where dL and dR are each side's
distance from the face midline. Frontal ~ 0, profile -> +/-1. It degrades
gracefully because it averages many pairs rather than trusting any one point,
and it is scale- and translation-free.

It is still an estimate. So this writes pose_review.html with the estimate
PRE-FILLED as a proposal, ordered by confidence, for you to correct — and the
correlation is only run against your corrected labels. If the estimate turns
out to disagree with you often, that is the finding, and the automatic number
gets dropped rather than reported.

Bins are coarse on purpose: frontal / three-quarter / profile / down-or-hidden.
Finer than that and the labelling is slower and less reliable, and the question
("does pose explain the misses") does not need degrees.
"""
import sys, glob, json, os, base64

# sae_embed MUST be imported before cv2: OpenCV's DNN module loads the system
# libonnxruntime, which then shadows the newer one this module links against and
# the import fails on a missing symbol version. Order matters, so do not tidy
# these into alphabetical order.
sys.path.insert(0, "/home/dtourolle/Development/Jray-project/scene-actor-extraction/build-ort")
import sae_embed

import numpy as np
import cv2

ROOT = "/home/dtourolle/Development/Jray-project/scene-actor-extraction/"
M = ROOT + "models/"
MESH = "/home/dtourolle/Development/rPPG-kahn/models/face_landmark.tflite"
CLIPS = ["5157344", "5157339"]
BINS = ["frontal", "three-quarter", "profile", "down-or-hidden"]

# Symmetric vertex pairs (subject-left, subject-right) on the MediaPipe mesh:
# outer eye corners, inner eye corners, cheeks, mouth corners, jaw.
PAIRS = [(33, 263), (133, 362), (130, 359), (243, 463),
         (61, 291), (91, 321), (146, 375), (58, 288), (172, 397), (215, 435)]
MIDLINE = [10, 168, 1, 4, 5, 195, 197, 152]      # forehead -> nose -> chin

net = cv2.dnn.readNetFromTFLite(MESH)
NAMES = net.getUnconnectedOutLayersNames()
LMI, PRI = NAMES.index("conv2d_21"), NAMES.index("conv2d_31")

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)


def mesh_pose(img, bbox, expand=1.6):
    """(yaw_ratio, presence) or (nan, 0). yaw_ratio in [-1, 1], 0 = frontal."""
    x, y, w, h = bbox
    cx, cy, s = x + w / 2, y + h / 2, max(w, h) * expand
    crop = cv2.getRectSubPix(img, (int(s), int(s)), (float(cx), float(cy)))
    net.setInput(cv2.dnn.blobFromImage(crop, 1 / 255.0, (192, 192), (0, 0, 0), swapRB=True))
    o = net.forward(NAMES)
    pres = 1 / (1 + np.exp(-float(o[PRI].ravel()[0])))
    lm = o[LMI].reshape(468, 3)[:, :2]
    mid = lm[MIDLINE]
    # least-squares midline direction, then signed distance of each pair member
    c = mid.mean(axis=0)
    u, _, _ = np.linalg.svd(mid - c)
    d = (mid - c)
    axis = np.linalg.svd(d.T @ d)[0][:, 0]        # principal direction of the midline
    normal = np.array([-axis[1], axis[0]])
    ratios = []
    for a, b in PAIRS:
        dl = float(np.dot(lm[a] - c, normal))
        dr = float(np.dot(lm[b] - c, normal))
        if abs(dl) + abs(dr) < 1e-6:
            continue
        ratios.append((abs(dl) - abs(dr)) / (abs(dl) + abs(dr)))
    return (float(np.median(ratios)) if ratios else np.nan), pres


def b64(img, size, q=72):
    ok, buf = cv2.imencode(".jpg", cv2.resize(img, (size, size)),
                           [cv2.IMWRITE_JPEG_QUALITY, q])
    return "data:image/jpeg;base64," + base64.b64encode(buf).decode() if ok else ""


items = []
for clip in CLIPS:
    lab = {os.path.basename(p): (os.path.basename(os.path.dirname(p)), 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"))}
    for frame in sorted({man[f]["frame"] for f in lab}):
        img = cv2.imread(f"frames/d{clip}_{frame}.png")
        dets = eng.detect(img)
        for fname, (person, path) in lab.items():
            m = man[fname]
            if m["frame"] != frame or m["idx"] >= len(dets):
                continue
            d = dets[m["idx"]]
            lm5 = np.array(d.landmarks, dtype=np.float32).reshape(5, 2)
            crop = sae_embed.align_face(img, lm5)
            if crop is None:
                continue
            yaw, pres = mesh_pose(img, d.bbox)
            a = abs(yaw) if not np.isnan(yaw) else 1.0
            guess = ("frontal" if a < 0.15 else "three-quarter" if a < 0.45
                     else "profile")
            if pres < 0.5:
                guess = "down-or-hidden"      # mesh could not fit at all
            ctx = cv2.imread(path)
            items.append({"file": fname, "clip": clip, "person": person,
                          "px": int(m["px"]), "yaw": None if np.isnan(yaw) else round(yaw, 3),
                          "pres": round(pres, 3), "guess": guess,
                          "ctx": b64(ctx, 140) if ctx is not None else "",
                          "ali": b64(np.asarray(crop), 112)})

# least-confident first: near a bin boundary, or the mesh could not fit
def uncertainty(it):
    if it["pres"] < 0.5:
        return 0.0
    a = abs(it["yaw"]) if it["yaw"] is not None else 1.0
    return min(abs(a - 0.15), abs(a - 0.45))
items.sort(key=uncertainty)

payload = json.dumps({"bins": BINS, "items": items})

HTML = """<meta charset="utf-8"><title>JRay — head pose labelling</title>
<style>
:root{color-scheme:dark}
body{margin:0;background:#14161a;color:#e6e8ea;font:14px/1.5 system-ui,sans-serif}
header{position:sticky;top:0;background:#181b20;border-bottom:1px solid #262b33;
 padding:12px 18px;display:flex;gap:16px;align-items:center;flex-wrap:wrap;z-index:5}
h1{font-size:15px;margin:0}
button{background:#232830;color:#e6e8ea;border:1px solid #262b33;border-radius:6px;
 padding:7px 12px;cursor:pointer;font:inherit}
button:hover{background:#2c323c}
button.go{background:#2f5d43;border-color:#3c7555}
.g{display:grid;grid-template-columns:repeat(auto-fill,minmax(300px,1fr));gap:10px;padding:14px}
.c{border:1px solid #262b33;border-radius:8px;padding:9px;display:flex;gap:9px;align-items:center}
.c.edited{border-color:#3c7555}
img{border-radius:5px;background:#000;display:block}
.m{color:#8b929c;font-size:11px}
.b{display:flex;flex-wrap:wrap;gap:4px;margin-top:5px}
.b button{padding:3px 7px;font-size:11px}
.b button.sel{background:#2f5d43;border-color:#3c7555}
</style>
<header><h1>Head pose</h1><span class="m" id="stat"></span>
<button class="go" id="exp">Export pose_labels.json</button></header>
<div class="g" id="g"></div>
<script>
const D=__PAYLOAD__; const pick={};
function render(){
 const g=document.getElementById('g'); g.innerHTML='';
 for(const it of D.items){
  const cur=pick[it.file]||it.guess;
  const c=document.createElement('div');
  c.className='c'+(pick[it.file]&&pick[it.file]!==it.guess?' edited':'');
  const b=D.bins.map(x=>`<button class="${x===cur?'sel':''}" data-f="${it.file}" data-b="${x}">${x}</button>`).join('');
  c.innerHTML=`<img src="${it.ctx}" width="88" height="88"><img src="${it.ali}" width="66" height="66">
   <div><div class="m">${it.person} · ${it.clip.slice(-3)} · ${it.px}px</div>
   <div class="m">yaw ${it.yaw===null?'—':it.yaw} · presence ${it.pres}</div>
   <div class="b">${b}</div></div>`;
  g.appendChild(c);
 }
 g.onclick=e=>{const t=e.target; if(t.dataset&&t.dataset.b){pick[t.dataset.f]=t.dataset.b; render();}};
 const ed=Object.entries(pick).filter(([f,v])=>v!==(D.items.find(i=>i.file===f)||{}).guess).length;
 document.getElementById('stat').textContent=`${D.items.length} crops · ${ed} corrections`;
}
document.getElementById('exp').onclick=()=>{
 const out={}; for(const it of D.items) out[it.file]={pose:pick[it.file]||it.guess,
   guess:it.guess, yaw:it.yaw, pres:it.pres, person:it.person, clip:it.clip};
 const a=document.createElement('a');
 a.href=URL.createObjectURL(new Blob([JSON.stringify(out,null,1)],{type:'application/json'}));
 a.download='pose_labels.json'; a.click();
};
render();
</script>
"""
out = "labelling/pose_review.html"
open(out, "w").write(HTML.replace("__PAYLOAD__", payload))
from collections import Counter
print(f"{out}  {os.path.getsize(out)/1e6:.1f} MB  {len(items)} crops", file=sys.stderr)
print(f"estimate: {dict(Counter(i['guess'] for i in items))}", file=sys.stderr)
print("\nestimated pose per person (does this match what you see?):", file=sys.stderr)
for p in sorted({i["person"] for i in items}):
    for clip in CLIPS:
        sub = [i for i in items if i["person"] == p and i["clip"] == clip]
        if sub:
            print(f"   {p} {clip[-3:]}: {dict(Counter(i['guess'] for i in sub))}",
                  file=sys.stderr)
print(f"\nopen file://{os.path.abspath(out)}", file=sys.stderr)