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

Usage:
  python scripts/movienet_eval.py \
    --gallery gallery_r50.h5 \
    --arcface models/arcface_w600k_r50.onnx \
    --gt eval/gt.json \
    --output eval/predictions_r50.json \
    [--build-dir build]

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 sys
from pathlib import Path

import numpy as np

sys.path.insert(0, str(Path(__file__).resolve().parent))
from sae_embed_loader import load_embedder
from sae_gallery import load_gallery_hdf5


def load_gallery(path: str) -> dict[str, dict]:
    """Return {imdb_id: {"name": str, "refs": np.ndarray[n_refs, dim]}}."""
    data = load_gallery_hdf5(Path(path))
    return {a["imdb_id"]: {"name": a["name"],
                           "refs": np.asarray(a["embeddings"], dtype=np.float32)}
            for a in data["actors"]}


def match(embedding: list[float], gallery: dict[str, dict]) -> tuple[str, float, dict[str, float]]:
    """Return (best_imdb_id, best_similarity, {imdb_id: similarity}).

    Similarity to an actor is the max dot product over that actor's reference
    embeddings; vectorised with numpy so a whole-library gallery stays fast.
    """
    vec = np.asarray(embedding, dtype=np.float32)
    scores: dict[str, float] = {
        imdb_id: float((actor["refs"] @ vec).max())
        for imdb_id, actor in gallery.items()
    }
    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("--build-dir", default="build",
                   help="Build directory containing the sae_embed module (default: build)")
    p.add_argument("--models-dir", default="models",
                   help="Directory containing ONNX models (default: models/)")
    args = p.parse_args()

    embedder = load_embedder(args.build_dir, args.models_dir, args.arcface)

    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)

    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)

    embed_results = [embedder.embed(str(p)) if p.exists() else None for p in crop_paths]

    predictions = []
    n_det_fail = 0
    n_correct = 0

    for entry, result in zip(gt_entries, embed_results):
        detection_failed = result is None or not result.ok
        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()