#!/usr/bin/env python3
"""
movienet_eval.py — embed probe crops and match against a gallery.

Usage:
  python scripts/movienet_eval.py \
    --gallery gallery_r50.json \
    --arcface models/arcface_w600k_r50.onnx \
    --gt eval/gt.json \
    --output eval/predictions_r50.json \
    [--yunet models/face_detection_yunet_2023mar.onnx] \
    [--embed-bin build/embed_faces]

Input  (--gt): list of {"crop": <path>, "imdb_id": <str>, "actor_name": <str>}
Output: list of {"crop", "gt", "pred", "similarity", "detection_failed", "all_scores"}
"""

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


def load_gallery(path: str) -> dict[str, dict]:
    """Return {imdb_id: {"name": str, "embeddings": [[float]]}}."""
    with open(path) as f:
        data = json.load(f)
    return {a["imdb_id"]: {"name": a["name"], "embeddings": a["embeddings"]}
            for a in data["actors"]}


def dot(a: list[float], b: list[float]) -> float:
    return sum(x * y for x, y in zip(a, b))


def embed_images(paths: list[Path], embed_bin: str, yunet: str, arcface: str) -> list[dict | None]:
    if not paths:
        return []
    cmd = [embed_bin, "--yunet", yunet, "--arcface", arcface] + [str(p) for p in paths]
    try:
        proc = subprocess.run(cmd, capture_output=True, text=True, check=True)
    except subprocess.CalledProcessError as e:
        print(f"[error] embed_faces failed:\n{e.stderr}", file=sys.stderr)
        return [None] * len(paths)
    try:
        return json.loads(proc.stdout)
    except json.JSONDecodeError as e:
        print(f"[error] embed_faces JSON parse error: {e}", file=sys.stderr)
        return [None] * len(paths)


def match(embedding: list[float], gallery: dict[str, dict]) -> tuple[str, float, dict[str, float]]:
    """Return (best_imdb_id, best_similarity, {imdb_id: similarity})."""
    scores: dict[str, float] = {}
    for imdb_id, actor in gallery.items():
        # max similarity across all reference embeddings for this actor
        scores[imdb_id] = max(dot(embedding, ref) for ref in actor["embeddings"])
    best_id = max(scores, key=lambda k: scores[k])
    return best_id, scores[best_id], scores


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--gallery",   required=True)
    p.add_argument("--arcface",   required=True)
    p.add_argument("--gt",        required=True)
    p.add_argument("--output",    required=True)
    p.add_argument("--yunet",     default="models/face_detection_yunet_2023mar.onnx")
    p.add_argument("--embed-bin", default="build/embed_faces")
    args = p.parse_args()

    gallery = load_gallery(args.gallery)
    print(f"[eval] gallery: {len(gallery)} actors", file=sys.stderr)

    with open(args.gt) as f:
        gt_entries = json.load(f)
    print(f"[eval] probe crops: {len(gt_entries)}", file=sys.stderr)

    # Batch all crops in one embed_faces call to amortise startup cost
    crop_paths = [Path(e["crop"]) for e in gt_entries]
    missing = [p for p in crop_paths if not p.exists()]
    if missing:
        print(f"[warn] {len(missing)} crop(s) not found on disk, skipping", file=sys.stderr)

    results_raw = embed_images(
        [p for p in crop_paths if p.exists()],
        args.embed_bin, args.yunet, args.arcface
    )

    # Re-index results back to original list (missing files get None)
    raw_iter = iter(results_raw)
    embed_results: list[dict | None] = []
    for p in crop_paths:
        embed_results.append(next(raw_iter) if p.exists() else None)

    predictions = []
    n_det_fail = 0
    n_correct = 0

    for entry, result in zip(gt_entries, embed_results):
        detection_failed = result is None or result.get("embedding") is None
        if detection_failed:
            n_det_fail += 1
            predictions.append({
                "crop": entry["crop"],
                "gt": entry["imdb_id"],
                "pred": None,
                "similarity": None,
                "detection_failed": True,
                "all_scores": {},
            })
            continue

        pred_id, sim, all_scores = match(result["embedding"], gallery)
        correct = pred_id == entry["imdb_id"]
        if correct:
            n_correct += 1

        predictions.append({
            "crop": entry["crop"],
            "gt": entry["imdb_id"],
            "pred": pred_id,
            "similarity": sim,
            "detection_failed": False,
            "all_scores": all_scores,
        })

    n_total = len(gt_entries)
    n_evaluated = n_total - n_det_fail
    rank1 = n_correct / n_evaluated * 100 if n_evaluated else 0
    print(f"[eval] detection failures: {n_det_fail}/{n_total}", file=sys.stderr)
    print(f"[eval] rank-1 accuracy:    {rank1:.1f}% ({n_correct}/{n_evaluated})", file=sys.stderr)

    out_path = Path(args.output)
    out_path.parent.mkdir(parents=True, exist_ok=True)
    with open(out_path, "w") as f:
        json.dump(predictions, f, indent=2)
    print(f"[eval] written → {out_path}", file=sys.stderr)


if __name__ == "__main__":
    main()