#!/usr/bin/env python3
"""Build labelling/review.html — a local page for correcting the labels.

One row per crop, ordered most-suspicious first:

    left    the person it is currently filed under (medoid of that person's
            hand-sorted crops, so the reference is one you trust)
    centre  the crop under review — context with the detection boxed, and
            beneath it the 112x112 the embedder actually receives
    right   the person it matches better, if any, with both probabilities

Pick a destination per row, then Export to download corrections.json and apply
it with apply_corrections.py. Nothing is moved by this script.

Self-contained: images are inlined as data URIs and the page is opened from
disk, so no server runs and no face crop leaves the machine.

Ordering is by P(other) - P(self), both from the global gallery sigmoid, so
rows where the evidence disagrees with the label float to the top and the
agreement cases sink. It is a review order, not a verdict — you are the
arbiter, which is the whole point of labelling by hand.
"""
import sys, glob, json, os, base64
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"
REF_CLIP = "5157344"          # the clip sorted by hand — reference faces come from here
CLIPS = ["5157344", "5157339"]

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


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


rows = []
for clip in CLIPS:
    man = {m["file"]: m for m in json.load(open(f"labelling/{clip}/manifest.json"))}
    placed = {os.path.basename(p): (os.path.basename(os.path.dirname(p)), p)
              for p in glob.glob(f"labelling/{clip}/*/*.jpg")}
    by_frame = {}
    for fname, (person, path) in placed.items():
        if fname in man and person != "unsorted":
            by_frame.setdefault(man[fname]["frame"], []).append((fname, person, path))
    for frame, items in sorted(by_frame.items()):
        img = cv2.imread(f"frames/d{clip}_{frame}.png")
        if img is None:
            continue
        dets = eng.detect(img)
        for fname, person, path in items:
            i = man[fname]["idx"]
            if i >= len(dets):
                continue
            lm = np.array(dets[i].landmarks, dtype=np.float32).reshape(5, 2)
            crop = sae_embed.align_face(img, lm)
            if crop is None:
                continue
            rows.append({"clip": clip, "person": person, "file": fname, "path": path,
                         "px": man[fname]["px"], "aligned": np.asarray(crop),
                         "emb": np.asarray(eng.embed_crop(crop), dtype=np.float32)})

people = sorted({r["person"] for r in rows})
E = np.stack([r["emb"] 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)

# reference face per person: medoid of their REF_CLIP crops
ref_img = {}
for k, p in enumerate(people):
    idx = [i for i in np.where(lab == k)[0] if rows[i]["clip"] == REF_CLIP]
    if not idx:
        idx = list(np.where(lab == k)[0])
    if not idx:
        continue
    sub = S[np.ix_(idx, idx)].copy()
    medoid = idx[int(np.argmax(sub.mean(axis=1)))]
    ref_img[p] = b64(rows[medoid]["aligned"], 112)

items = []
for i, r in enumerate(rows):
    k = lab[i]
    same = [j for j in np.where(lab == k)[0] if j != i]
    p_self = cal.probability(float(S[i, same].max())) if same else 0.0
    best_other, p_other = None, 0.0
    for k2, p2 in enumerate(people):
        if k2 == k:
            continue
        other = np.where(lab == k2)[0]
        if not len(other):
            continue
        pv = cal.probability(float(S[i, other].max()))
        if pv > p_other:
            p_other, best_other = pv, p2
    ctx = cv2.imread(r["path"])
    items.append({
        "file": r["file"], "clip": r["clip"], "person": r["person"],
        "px": int(r["px"]), "p_self": round(p_self, 3), "p_other": round(p_other, 3),
        "other": best_other, "delta": round(p_other - p_self, 3),
        "ctx": b64(ctx, 150) if ctx is not None else "",
        "ali": b64(r["aligned"], 112),
    })
items.sort(key=lambda x: -x["delta"])

payload = json.dumps({"people": people, "refs": ref_img, "items": items})

HTML = """<meta charset="utf-8"><title>JRay — label review</title>
<style>
:root{color-scheme:dark;--bg:#14161a;--fg:#e6e8ea;--mut:#8b929c;--line:#262b33;--warn:#e0654a;--ok:#4a9d6a}
body{margin:0;background:var(--bg);color:var(--fg);font:14px/1.5 system-ui,sans-serif}
header{position:sticky;top:0;background:#181b20;border-bottom:1px solid var(--line);
 padding:12px 18px;display:flex;gap:18px;align-items:center;flex-wrap:wrap;z-index:5}
h1{font-size:15px;margin:0;font-weight:600}
.stat{color:var(--mut);font-size:13px}
button{background:#232830;color:var(--fg);border:1px solid var(--line);border-radius:6px;
 padding:7px 13px;cursor:pointer;font:inherit}
button:hover{background:#2c323c}
button.go{background:#2f5d43;border-color:#3c7555}
.row{display:grid;grid-template-columns:150px 1fr 190px;gap:20px;align-items:center;
 padding:14px 18px;border-bottom:1px solid var(--line)}
.row.flag{background:#1e1719}
.row.done{opacity:.4}
.cell{display:flex;gap:10px;align-items:center}
img{border-radius:5px;display:block;background:#000}
.lab{font-weight:600;font-size:15px}
.mut{color:var(--mut);font-size:12px}
.p{font-variant-numeric:tabular-nums}
.hi{color:var(--warn);font-weight:600}
.choices{display:flex;flex-wrap:wrap;gap:6px}
.choices button{padding:5px 10px;font-size:13px}
.choices button.sel{background:#2f5d43;border-color:#3c7555}
.legend{padding:10px 18px;color:var(--mut);font-size:12px;border-bottom:1px solid var(--line)}
</style>
<header>
  <h1>Label review</h1>
  <span class="stat" id="stat"></span>
  <button id="exp" class="go">Export corrections.json</button>
  <button id="onlyflag">Show only disagreements</button>
</header>
<div class="legend">Left: the person this crop is filed under. Centre: the crop (context with the
detection boxed, and the 112&times;112 the embedder actually sees). Right: the person it matches
better, if any. Ordered by P(other) &minus; P(self) &mdash; disagreements first.</div>
<div id="list"></div>
<script>
const D = __PAYLOAD__;
const choice = {};
const list = document.getElementById('list');

function render(){
  list.innerHTML = '';
  const flagOnly = document.body.dataset.flag === '1';
  for (const it of D.items){
    if (flagOnly && it.delta <= 0) continue;
    const row = document.createElement('div');
    row.className = 'row' + (it.delta > 0 ? ' flag' : '') + (choice[it.file] ? ' done' : '');

    const left = document.createElement('div');
    left.className = 'cell';
    left.innerHTML = `<img src="${D.refs[it.person]||''}" width="72" height="72">
      <div><div class="lab">${it.person}</div>
      <div class="mut p">P(self) ${it.p_self.toFixed(3)}</div></div>`;

    const mid = document.createElement('div');
    mid.className = 'cell';
    mid.innerHTML = `<img src="${it.ctx}" width="120" height="120">
      <img src="${it.ali}" width="90" height="90">
      <div><div class="mut">${it.clip} &middot; ${it.px}px</div>
      <div class="mut">${it.file}</div></div>`;

    const right = document.createElement('div');
    const worse = it.delta > 0;
    right.innerHTML = it.other
      ? `<div class="cell"><img src="${D.refs[it.other]||''}" width="56" height="56">
         <div><div class="lab ${worse?'hi':''}">${it.other}</div>
         <div class="mut p ${worse?'hi':''}">P ${it.p_other.toFixed(3)}</div></div></div>`
      : '<div class="mut">—</div>';

    const ch = document.createElement('div');
    ch.className = 'choices';
    for (const p of D.people.concat(['discard'])){
      const b = document.createElement('button');
      b.textContent = p === it.person ? p + ' (keep)' : p;
      if (choice[it.file] === p || (!choice[it.file] && p === it.person)) b.classList.add('sel');
      b.onclick = () => { choice[it.file] = p; render(); };
      ch.appendChild(b);
    }
    right.appendChild(ch);

    row.append(left, mid, right);
    list.appendChild(row);
  }
  const changed = Object.entries(choice).filter(([f,p]) =>
    p !== (D.items.find(i=>i.file===f)||{}).person).length;
  document.getElementById('stat').textContent =
    `${D.items.length} crops · ${D.items.filter(i=>i.delta>0).length} disagreements · ${changed} changes staged`;
}

document.getElementById('onlyflag').onclick = () => {
  document.body.dataset.flag = document.body.dataset.flag === '1' ? '0' : '1';
  render();
};
document.getElementById('exp').onclick = () => {
  const out = {};
  for (const it of D.items){
    const p = choice[it.file] || it.person;
    if (p !== it.person) out[it.file] = {from: it.person, to: p, clip: it.clip};
  }
  const blob = new Blob([JSON.stringify(out, null, 1)], {type:'application/json'});
  const a = document.createElement('a');
  a.href = URL.createObjectURL(blob); a.download = 'corrections.json'; a.click();
};
render();
</script>
"""

os.makedirs("labelling", exist_ok=True)
out = "labelling/review.html"
with open(out, "w") as f:
    f.write(HTML.replace("__PAYLOAD__", payload))
size = os.path.getsize(out) / 1e6
flagged = sum(1 for i in items if i["delta"] > 0)
print(f"{out}  {size:.1f} MB  {len(items)} crops, {flagged} disagreements", file=sys.stderr)
print(f"open file://{os.path.abspath(out)}", file=sys.stderr)