#!/usr/bin/env python3
"""
calibration_chart.py — plot each model's calibrated P(match|similarity) sigmoid,
from the (a, b) fitted into each gallery's HDF5 /calibration group. Shows
discriminative power: a steeper curve (larger |a|) separates positive/negative
pairs more sharply at the same decision boundary.

Usage:
  python scripts/docs/calibration_chart.py --out docs/assets/images/calibration_curves.png
"""
from __future__ import annotations

import argparse
from pathlib import Path

import h5py
import numpy as np
import matplotlib.pyplot as plt

MODELS = [
    ("arcface_w600k_r50", "ArcFace w600k-R50"),
    ("arcface_r18", "ArcFace R18"),
    ("arcface_w600k_mbf", "ArcFace w600k-MBF"),
    ("LVFace-B_Glint360K", "LVFace-B Glint360K"),
]
COLOURS = ["#2a78d6", "#008300", "#e87ba4", "#eda100"]

REPO = Path(__file__).resolve().parent.parent.parent


def load_calibrations() -> list[dict]:
    out = []
    for slug, label in MODELS:
        path = REPO / f"experiments/galleries/gallery_{slug}.h5"
        if not path.exists():
            print(f"[calibration_chart] skip {slug}: gallery not found at {path}")
            continue
        with h5py.File(path, "r") as f:
            if "calibration" not in f:
                print(f"[calibration_chart] skip {slug}: no calibration in gallery "
                     f"(run a replay against it once to fit and embed one)")
                continue
            cal = f["calibration"]
            out.append({"slug": slug, "label": label,
                       "a": float(cal.attrs["a"]), "b": float(cal.attrs["b"])})
    return out


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--out", required=True)
    args = p.parse_args()

    models = load_calibrations()
    if not models:
        raise SystemExit("no galleries had embedded calibration — run a replay "
                         "against each gallery once first (see identity_matcher_node.hpp)")

    sim = np.linspace(-1, 1, 400)
    fig, ax = plt.subplots(figsize=(7.5, 4.8), dpi=150)

    for m, colour in zip(models, COLOURS):
        p_match = 1.0 / (1.0 + np.exp(-(m["a"] * sim + m["b"])))
        boundary = -m["b"] / m["a"]
        ax.plot(sim, p_match, color=colour, linewidth=2,
               label=f"{m['label']}  (a={m['a']:.1f}, boundary@P=0.5: sim={boundary:.2f})")

    ax.axhline(0.5, color="#999999", linewidth=1, linestyle="--", zorder=0)
    ax.set_xlabel("cosine similarity")
    ax.set_ylabel("P(match)")
    ax.set_title("Calibrated P(match | similarity), per embedding model")
    ax.set_xlim(-1, 1)
    ax.set_ylim(0, 1)
    ax.legend(loc="upper left", fontsize=8, frameon=False)
    ax.spines["top"].set_visible(False)
    ax.spines["right"].set_visible(False)
    fig.tight_layout()

    out_path = Path(args.out)
    out_path.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(out_path)
    print(f"[calibration_chart] wrote {out_path} ({len(models)} models)")
    for m in models:
        print(f"  {m['label']}: a={m['a']:.2f} b={m['b']:.2f} "
              f"boundary(P=0.5)=sim{-m['b']/m['a']:.3f}")


if __name__ == "__main__":
    main()