#!/usr/bin/env python3
"""
dump_scene_montage.py — one BEST and one WORST frame per X-Ray scene, split into
onscreen vs. offscreen actor identification (TPI / FPI / FN).

For each X-Ray scene (scenes.csv span), scores every sampled second by a simple
per-second Jaccard agreement (TPI / (TPI+FPI+FN), same spirit as second_score.py)
and picks the single best-agreement and single worst-agreement second. Each gets
one output frame: the full frame (not a face crop) with a solid box drawn for
every currently-active TPI/FPI actor who has a REAL detection backing them, plus a
black caption panel below with two columns — Onscreen (has a real detection) and
Offscreen (no real detection: FN misses, and "ghost" detections where the tracker
is re-emitting a frozen last-known bbox with nothing there — see
docs/model-bakeoff.md) — names colour-coded by bucket, with a legend.

A predicted bbox is checked against the dump's OWN raw per-frame face detections
(IoU) to tell a real detection from a ghost. Ghosts are NEVER drawn as boxes (they
have no real screen position); they only appear as a name in the Offscreen column.

Frames where at least one FPI name isn't in the film's cast AT ALL (an out-of-cast
misID, not just a right-actor/wrong-scene timing slip) are also copied into
<out-dir>/out_of_cast_fpi/ for quick review of the most confident wrong answers.

Requires the raw per-frame annotations from `replay.py --raw-out` (bboxes aren't
in the merged pred.json) and the film's HDF5 dump (for ghost-checking against real
detections).

Usage:
  python scripts/optimizer/dump_scene_montage.py \
      --raw raw.jsonl --dump experiments/dumps/.../dump_X.h5 \
      --xray experiments/xray/.../900_The_Many_Saints_Of_Newark \
      --movie "/mnt/movies/.../X.mp4" \
      --gallery experiments/galleries/gallery_LVFace-B_Glint360K.h5 \
      --out-dir experiments/results/holdout/montage/many_saints --scene 5
"""
from __future__ import annotations

import argparse
import csv
import json
import subprocess
import sys
from pathlib import Path

import cv2
import h5py
import numpy as np

REPO = Path(__file__).resolve().parent.parent.parent
sys.path.insert(0, str(REPO / "scripts" / "optimizer"))
sys.path.insert(0, str(REPO / "scripts" / "validation"))

from sample_eval import load_gallery_keys  # noqa: E402
from identity import keys_for  # noqa: E402

COLOUR_TPI = (60, 200, 0)     # green, BGR
COLOUR_FPI = (0, 60, 220)     # red, BGR
CAPTION_H = 28  # px per line in the bottom strip


def load_scene_spans(xray_dir: str):
    """scene_id -> (t0_sec, t1_sec), from scenes.csv (ms)."""
    spans = {}
    with open(Path(xray_dir) / "scenes.csv", newline="", encoding="utf-8") as f:
        for r in csv.DictReader(f):
            sn = (r.get("scene") or "").strip()
            try:
                spans[sn] = (float(r["start"]) / 1000.0, float(r["end"]) / 1000.0)
            except (KeyError, ValueError):
                continue
    return spans


def load_film_cast(xray_dir: str) -> set:
    """Every actor key X-Ray credits ANYWHERE in the film — used to tell an
    out-of-cast misID (named someone who isn't even in this film) apart from an
    in-cast timing slip (right actor, wrong scene), same distinction as
    second_score.py's FPI_misid vs FPI_incast."""
    keys = set()
    with open(Path(xray_dir) / "people.csv", newline="", encoding="utf-8") as f:
        for r in csv.DictReader(f):
            nm = (r.get("name_id") or "").strip()
            person = (r.get("person") or "").strip()
            if nm or person:
                keys |= keys_for(imdb_id=nm, name=person)
    return keys


def load_scene_cast(xray_dir: str):
    """scene_id -> set of actor key-frozensets X-Ray lists as present."""
    id_to_name = {}
    with open(Path(xray_dir) / "people.csv", newline="", encoding="utf-8") as f:
        for r in csv.DictReader(f):
            nm = (r.get("name_id") or "").strip()
            if nm:
                id_to_name[nm] = (r.get("person") or "").strip()

    scene_cast: dict[str, set] = {}
    with open(Path(xray_dir) / "people_in_scenes.csv", newline="", encoding="utf-8") as f:
        for r in csv.DictReader(f):
            sn = (r.get("scene") or "").strip()
            nm = (r.get("name_id") or "").strip()
            if sn and nm:
                scene_cast.setdefault(sn, set()).add(
                    frozenset(keys_for(imdb_id=nm, name=id_to_name.get(nm))))
    return scene_cast


def load_raw_by_second(raw_path: str):
    by_second: dict[int, list] = {}
    with open(raw_path) as f:
        for line in f:
            sa = json.loads(line)
            if sa.get("eof"):
                continue
            by_second[int(sa["timestamp_sec"])] = sa.get("visible_actors", [])
    return by_second


def load_dump_faces_by_second(dump_path: str):
    """second (int) -> list of raw detected bboxes (x,y,w,h), for ghost-checking.
    A predicted actor's bbox is real iff it overlaps one of these; a bbox with no
    overlap at all is a frozen/stale re-emission, not an actual detection."""
    by_second: dict[int, list] = {}
    with h5py.File(dump_path, "r") as f:
        ts = f["frames/timestamp_sec"][:]
        off = f["frames/face_offset"][:]
        cnt = f["frames/face_count"][:]
        bbox = f["faces/bbox"][:]
        for i in range(len(ts)):
            s, n = int(off[i]), int(cnt[i])
            by_second[int(ts[i])] = [tuple(b) for b in bbox[s:s + n]]
    return by_second


def iou(a, b):
    ax, ay, aw, ah = a
    bx, by, bw, bh = b
    ix0, iy0 = max(ax, bx), max(ay, by)
    ix1, iy1 = min(ax + aw, bx + bw), min(ay + ah, by + bh)
    iw, ih = max(0.0, ix1 - ix0), max(0.0, iy1 - iy0)
    inter = iw * ih
    union = aw * ah + bw * bh - inter
    return inter / union if union > 0 else 0.0


def is_ghost(bbox, real_boxes, iou_thresh=0.3):
    return not any(iou(bbox, rb) >= iou_thresh for rb in real_boxes)


def actor_key(a: dict) -> frozenset:
    return frozenset(keys_for(imdb_id=a.get("imdb_id"), tmdb_id=a.get("tmdb_id"),
                              jellyfin_id=a.get("jellyfin_id"), name=a.get("name")))


COLOUR_FN = (220, 130, 0)   # blue, BGR
LEGEND = (("TPI (correct)", COLOUR_TPI), ("FPI (wrong)", COLOUR_FPI),
         ("FN (missed)", COLOUR_FN))


def render_frame(frame_path: Path, t: int, tpi_boxes: list, fpi_boxes: list, entries: list):
    """entries: list of (name, bucket, onscreen) — bucket in {tpi,fpi,fn},
    onscreen=True iff a real detected face backs this name at this second. Ghost
    detections (bucket fpi/tpi but no real face — see is_ghost) are never drawn as
    boxes: they have no real screen position, they only ever appear in the
    Offscreen column."""
    img = cv2.imread(str(frame_path))
    if img is None:
        return None

    for name, bbox, sim in tpi_boxes:
        x, y, w, h = (int(v) for v in bbox)
        cv2.rectangle(img, (x, y), (x + w, y + h), COLOUR_TPI, 2)
        _label(img, (x, y), f"{name} {sim*100:.0f}%", COLOUR_TPI)
    for name, bbox, sim in fpi_boxes:
        x, y, w, h = (int(v) for v in bbox)
        cv2.rectangle(img, (x, y), (x + w, y + h), COLOUR_FPI, 2)
        _label(img, (x, y), f"{name} {sim*100:.0f}%", COLOUR_FPI)

    h_img, w_img = img.shape[:2]
    bucket_colour = {"tpi": COLOUR_TPI, "fpi": COLOUR_FPI, "fn": COLOUR_FN}
    onscreen = [(n, bucket_colour[b]) for n, b, on in entries if on]
    offscreen = [(n, bucket_colour[b]) for n, b, on in entries if not on]

    n_rows = max(len(onscreen), len(offscreen), 1)
    header_h = 24
    legend_h = CAPTION_H
    table_h = header_h + n_rows * CAPTION_H + legend_h + 16
    canvas = np.zeros((h_img + table_h, w_img, 3), dtype=np.uint8)  # black bg
    canvas[:h_img] = img

    col_x = (8, w_img // 2 + 8)
    cv2.putText(canvas, f"t={t}s", (8, 16), cv2.FONT_HERSHEY_SIMPLEX, 0.5,
                (255, 255, 255), 1, cv2.LINE_AA)
    y0 = h_img + header_h
    cv2.putText(canvas, "Onscreen", (col_x[0], y0), cv2.FONT_HERSHEY_SIMPLEX, 0.55,
                (255, 255, 255), 1, cv2.LINE_AA)
    cv2.putText(canvas, "Offscreen", (col_x[1], y0), cv2.FONT_HERSHEY_SIMPLEX, 0.55,
                (255, 255, 255), 1, cv2.LINE_AA)
    cv2.line(canvas, (col_x[1] - 8, h_img), (col_x[1] - 8, h_img + table_h),
             (90, 90, 90), 1)

    for i in range(n_rows):
        y = y0 + CAPTION_H * (i + 1)
        if i < len(onscreen):
            name, colour = onscreen[i]
            cv2.putText(canvas, name, (col_x[0], y), cv2.FONT_HERSHEY_SIMPLEX, 0.5,
                        colour, 1, cv2.LINE_AA)
        if i < len(offscreen):
            name, colour = offscreen[i]
            cv2.putText(canvas, name, (col_x[1], y), cv2.FONT_HERSHEY_SIMPLEX, 0.5,
                        colour, 1, cv2.LINE_AA)

    ly = y0 + CAPTION_H * (n_rows + 1) + 4
    lx = 8
    for label, colour in LEGEND:
        (tw, _), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.42, 1)
        cv2.rectangle(canvas, (lx, ly - 10), (lx + 12, ly + 2), colour, cv2.FILLED)
        cv2.putText(canvas, label, (lx + 18, ly), cv2.FONT_HERSHEY_SIMPLEX, 0.42,
                    (200, 200, 200), 1, cv2.LINE_AA)
        lx += tw + 40
    return canvas


def _label(img, pt, text, colour):
    x, y = pt
    (tw, th), _ = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
    strip_y0 = max(0, y - th - 4)
    cv2.rectangle(img, (x, strip_y0), (x + tw + 4, y), colour, cv2.FILLED)
    cv2.putText(img, text, (x + 2, y - 2), cv2.FONT_HERSHEY_SIMPLEX, 0.5,
                (255, 255, 255), 1, cv2.LINE_AA)


def extract_frame(movie: str, t: float, out_path: Path):
    out_path.parent.mkdir(parents=True, exist_ok=True)
    subprocess.run(
        ["ffmpeg", "-y", "-ss", str(t), "-i", movie, "-frames:v", "1",
         "-q:v", "2", str(out_path)],
        check=True, capture_output=True)


def classify_second(t: int, gt_cast: set, film_cast: set, raw_by_second: dict,
                    dump_faces_by_second: dict):
    """One second's TPI/FPI/FN classification: (score, tpi_boxes, fpi_boxes,
    entries, has_outofcast). score = Jaccard-style agreement in [0,1], used to
    rank seconds for best/worst picking."""
    actors = raw_by_second.get(t, [])
    real_boxes = dump_faces_by_second.get(t, [])

    tpi_boxes, fpi_boxes = [], []
    entries = []  # (name, bucket, onscreen)
    cur_state: dict[frozenset, str] = {}
    has_outofcast = False

    for a in actors:
        if a.get("actor_idx", -1) < 0:
            continue
        key = actor_key(a)
        name = a.get("name", "?")
        bbox = tuple(a["bbox"])
        sim = a.get("similarity", 0.0)
        ghost = is_ghost(bbox, real_boxes)
        hit = any(key & g for g in gt_cast)

        if ghost:
            cur_state[key] = "ghost"
            entries.append((name, "fpi" if not hit else "tpi", False))
        elif hit:
            tpi_boxes.append((name, bbox, sim))
            cur_state[key] = "tpi"
            entries.append((name, "tpi", True))
        else:
            fpi_boxes.append((name, bbox, sim))
            cur_state[key] = "fpi"
            entries.append((name, "fpi", True))

        if not hit and not (key & film_cast):
            has_outofcast = True  # named someone not in the film at all (FPI_misid)

    tpi_keys = [k for k, state in cur_state.items() if state in ("tpi", "ghost")]
    fn_count = 0
    for g in gt_cast:
        if any(g & k for k in tpi_keys):
            continue
        nm = next((x.split("name:", 1)[1] for x in g if x.startswith("name:")), None)
        entries.append((nm or next(iter(g), "?"), "fn", False))
        fn_count += 1

    tp = sum(1 for _, b, on in entries if b == "tpi" and on)
    fp = sum(1 for _, b, on in entries if b == "fpi")
    union = tp + fp + fn_count
    score = tp / union if union else 1.0  # both-empty = perfect agreement
    return score, tpi_boxes, fpi_boxes, entries, has_outofcast


def process_scene(scene_id: str, t0: float, t1: float, gt_cast: set, film_cast: set,
                  raw_by_second: dict, dump_faces_by_second: dict,
                  gallery_keys: set | None, movie: str, out_dir: Path,
                  outofcast_dir: Path):
    if gallery_keys is not None:
        gt_cast = {g for g in gt_cast if g & gallery_keys}

    per_second = {}
    for t in range(int(t0), int(t1)):
        per_second[t] = classify_second(t, gt_cast, film_cast, raw_by_second,
                                        dump_faces_by_second)
    if not per_second:
        return []

    best_t = max(per_second, key=lambda t: per_second[t][0])
    worst_t = min(per_second, key=lambda t: per_second[t][0])

    manifest = []
    for label, t in (("best", best_t), ("worst", worst_t)):
        score, tpi_boxes, fpi_boxes, entries, has_outofcast = per_second[t]
        fname = f"{scene_id}_{label}_t{t:06d}.jpg"
        out_path = out_dir / fname
        try:
            extract_frame(movie, t, out_path)
            canvas = render_frame(out_path, t, tpi_boxes, fpi_boxes, entries)
            if canvas is not None:
                cv2.imwrite(str(out_path), canvas)
            manifest.append({"label": label, "t": t, "score": round(score, 3),
                             "entries": entries, "file": fname,
                             "outofcast": has_outofcast})
            print(f"[scene {scene_id}] {label} t={t}s score={score:.2f} "
                  f"entries={entries}", file=sys.stderr)
            if has_outofcast:
                outofcast_dir.mkdir(parents=True, exist_ok=True)
                cv2.imwrite(str(outofcast_dir / fname), cv2.imread(str(out_path)))
        except subprocess.CalledProcessError as e:
            print(f"[dump_scene_montage] ffmpeg failed at t={t}: {e}", file=sys.stderr)

    return manifest


def main():
    p = argparse.ArgumentParser(description=__doc__,
                                formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("--raw", required=True, help="raw per-frame annotations (replay.py --raw-out)")
    p.add_argument("--dump", required=True, help="film's HDF5 embedding dump (for ghost-checking)")
    p.add_argument("--xray", required=True)
    p.add_argument("--movie", required=True)
    p.add_argument("--gallery")
    p.add_argument("--out-dir", required=True)
    p.add_argument("--scene", help="only process this X-Ray scene id (default: all)")
    args = p.parse_args()

    spans = load_scene_spans(args.xray)
    scene_cast = load_scene_cast(args.xray)
    film_cast = load_film_cast(args.xray)
    raw_by_second = load_raw_by_second(args.raw)
    dump_faces_by_second = load_dump_faces_by_second(args.dump)
    gk = load_gallery_keys(args.gallery) if args.gallery else None

    out_dir = Path(args.out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    outofcast_dir = out_dir / "out_of_cast_fpi"

    scene_ids = [args.scene] if args.scene else sorted(spans, key=lambda s: spans[s][0])
    all_manifest = {}
    for sn in scene_ids:
        if sn not in spans:
            print(f"[dump_scene_montage] unknown scene id: {sn}", file=sys.stderr)
            continue
        t0, t1 = spans[sn]
        gt_cast = scene_cast.get(sn, set())
        scene_dir = out_dir / f"scene_{sn}"
        scene_dir.mkdir(parents=True, exist_ok=True)
        m = process_scene(sn, t0, t1, gt_cast, film_cast, raw_by_second,
                          dump_faces_by_second, gk, args.movie, scene_dir, outofcast_dir)
        all_manifest[sn] = m

    (out_dir / "manifest.json").write_text(json.dumps(all_manifest, indent=2, ensure_ascii=False))
    total = sum(len(v) for v in all_manifest.values())
    n_outofcast = sum(1 for v in all_manifest.values() for r in v if r.get("outofcast"))
    print(f"[dump_scene_montage] wrote {total} best/worst frames across "
          f"{len(all_manifest)} scenes to {out_dir} "
          f"({n_outofcast} copied to {outofcast_dir})", file=sys.stderr)


if __name__ == "__main__":
    main()
