#!/usr/bin/env python3
"""
movienet_score.py — compare model predictions against ground truth.

Usage:
  python scripts/movienet_score.py \
    --gt eval/gt.json \
    --predictions eval/predictions_r50.json:R50:167MB \
                  eval/predictions_r18.json:R18:46MB \
                  eval/predictions_mbf.json:MBF:13MB \
    [--output eval/report.md]

Each --predictions value is  <path>:<label>:<size>  (size is display-only).
"""

import argparse
import json
import sys
from collections import defaultdict
from pathlib import Path


def load(path: str) -> list[dict]:
    with open(path) as f:
        return json.load(f)


def score(predictions: list[dict], gt_map: dict[str, str]) -> dict:
    n_total = len(predictions)
    n_det_fail = sum(1 for p in predictions if p["detection_failed"])
    evaluated = [p for p in predictions if not p["detection_failed"]]

    n_correct = sum(1 for p in evaluated if p["pred"] == p["gt"])
    rank1 = n_correct / len(evaluated) * 100 if evaluated else 0.0

    sims_correct = [p["similarity"] for p in evaluated if p["pred"] == p["gt"]]
    mean_sim = sum(sims_correct) / len(sims_correct) if sims_correct else 0.0

    # Per-actor recall
    per_actor: dict[str, dict] = defaultdict(lambda: {"correct": 0, "total": 0, "name": ""})
    for p in evaluated:
        actor_id = p["gt"]
        per_actor[actor_id]["total"] += 1
        per_actor[actor_id]["name"] = gt_map.get(actor_id, actor_id)
        if p["pred"] == actor_id:
            per_actor[actor_id]["correct"] += 1

    return {
        "n_total": n_total,
        "n_det_fail": n_det_fail,
        "n_evaluated": len(evaluated),
        "rank1": rank1,
        "mean_sim_correct": mean_sim,
        "per_actor": dict(per_actor),
    }


def render_table(rows: list[dict], headers: list[str]) -> str:
    col_widths = [max(len(h), max(len(str(r[h])) for r in rows)) for h in headers]
    sep = "| " + " | ".join("-" * w for w in col_widths) + " |"
    header = "| " + " | ".join(h.ljust(w) for h, w in zip(headers, col_widths)) + " |"
    lines = [header, sep]
    for r in rows:
        lines.append("| " + " | ".join(str(r[h]).ljust(w) for h, w in zip(headers, col_widths)) + " |")
    return "\n".join(lines)


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--gt",          required=True)
    p.add_argument("--predictions", required=True, nargs="+",
                   metavar="PATH:LABEL:SIZE")
    p.add_argument("--output",      default=None)
    args = p.parse_args()

    gt_entries = load(args.gt)
    gt_map = {e["imdb_id"]: e["actor_name"] for e in gt_entries}

    models = []
    for spec in args.predictions:
        parts = spec.split(":")
        if len(parts) != 3:
            print(f"[error] expected PATH:LABEL:SIZE, got: {spec}", file=sys.stderr)
            sys.exit(1)
        path, label, size = parts
        preds = load(path)
        s = score(preds, gt_map)
        models.append({"label": label, "size": size, "score": s})

    # ── Summary table ────────────────────────────────────────────────────────────
    summary_rows = []
    for m in models:
        s = m["score"]
        det_fail_pct = s["n_det_fail"] / s["n_total"] * 100 if s["n_total"] else 0
        summary_rows.append({
            "Model":    m["label"],
            "Rank-1":   f"{s['rank1']:.1f}%",
            "Det.Fail": f"{det_fail_pct:.1f}%",
            "Mean-sim": f"{s['mean_sim_correct']:.3f}",
            "Probes":   str(s["n_total"]),
            "Size":     m["size"],
        })

    summary_table = render_table(
        summary_rows,
        ["Model", "Rank-1", "Det.Fail", "Mean-sim", "Probes", "Size"]
    )

    # ── Per-actor table (using first model's actor list as reference) ────────────
    all_actor_ids = sorted({e["imdb_id"] for e in gt_entries})
    actor_rows = []
    for actor_id in all_actor_ids:
        row = {"Actor": gt_map.get(actor_id, actor_id)}
        for m in models:
            pa = m["score"]["per_actor"].get(actor_id, {"correct": 0, "total": 0})
            recall = pa["correct"] / pa["total"] * 100 if pa["total"] else 0.0
            row[m["label"]] = f"{recall:.0f}% ({pa['correct']}/{pa['total']})"
        actor_rows.append(row)

    actor_headers = ["Actor"] + [m["label"] for m in models]
    actor_table = render_table(actor_rows, actor_headers)

    # ── Assemble report ──────────────────────────────────────────────────────────
    report = f"""# MovieNet Validation Report

## Summary

{summary_table}

## Per-Actor Recall

{actor_table}
"""

    print(report)
    if args.output:
        out = Path(args.output)
        out.parent.mkdir(parents=True, exist_ok=True)
        out.write_text(report)
        print(f"[score] report written → {out}", file=sys.stderr)


if __name__ == "__main__":
    main()