#!/usr/bin/env python3
"""
experiment_charts.py — generate the rep4/held-out figures referenced by the docs,
from the experiment artifacts under experiments/ (no hardcoded numbers).

Figures:
  holdout_f1_by_film.png    — held-out per-film F1 vs. the training-set fit
  rep4_matrix_f1.png        — all 16 bake-off combos, colored by model
  de_search_landscape.png   — DE search space: prob_threshold x extinction_sec, F1 as color
  downton_ghost_timeline.png— detector face_count vs. tracker output through the credits

Usage:
  python scripts/docs/experiment_charts.py --out-dir docs/assets/images
"""
from __future__ import annotations

import argparse
import json
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
from matplotlib.colors import LinearSegmentedColormap

REPO = Path(__file__).resolve().parent.parent.parent
RESULTS = REPO / "experiments/results"

# Same model -> color mapping as calibration_chart.py, so identity is stable
# across every figure in the report.
MODEL_COLOURS = {
    "arcface_w600k_r50": ("ArcFace w600k-R50", "#2a78d6"),
    "arcface_r18": ("ArcFace R18", "#008300"),
    "arcface_w600k_mbf": ("ArcFace w600k-MBF", "#e87ba4"),
    "LVFace-B_Glint360K": ("LVFace-B Glint360K", "#eda100"),
}
INK = "#0b0b0b"
MUTED = "#898781"
GRID = "#e1e0d9"
SURFACE = "#fcfcfb"
BLUE = "#2a78d6"
GREEN = "#008300"
RED = "#e34948"

plt.rcParams.update({
    "figure.facecolor": SURFACE,
    "axes.facecolor": SURFACE,
    "savefig.facecolor": SURFACE,
    "text.color": INK,
    "axes.edgecolor": MUTED,
    "axes.labelcolor": INK,
    "xtick.color": MUTED,
    "ytick.color": MUTED,
    "axes.grid": True,
    "grid.color": GRID,
    "grid.linewidth": 0.8,
    "axes.spines.top": False,
    "axes.spines.right": False,
    "font.size": 11,
})


def training_best() -> dict:
    with open(RESULTS / "rep4_best_LVFace-B_Glint360K_full_exp.json") as f:
        return json.load(f)["best"]


def fig_holdout_f1(out: Path):
    with open(RESULTS / "holdout/holdout_scores.json") as f:
        films = json.load(f)["per_film"]
    films = sorted(films, key=lambda d: d["f1"])
    names = [d["name"] for d in films]
    f1 = [d["f1"] * 100 for d in films]
    train_f1 = training_best()["f1"] * 100
    macro = float(np.mean(f1))

    fig, ax = plt.subplots(figsize=(9, 4.2))
    ax.grid(axis="y", visible=False)
    bars = ax.barh(names, f1, height=0.55, color=BLUE, zorder=3)
    for b, v, d in zip(bars, f1, films):
        note = f"{v:.1f}%"
        if d["FPI_misid"]:
            note += f"   ({d['FPI_misid']} misIDs)"
        ax.text(v + 1, b.get_y() + b.get_height() / 2, note,
                va="center", ha="left", fontsize=10, color=INK)
    ax.axvline(train_f1, color=MUTED, lw=1.5, ls="--", zorder=2)
    ax.text(train_f1 + 0.7, len(names) - 0.35, f"training-set fit {train_f1:.1f}%",
            color=MUTED, fontsize=9.5, ha="left", va="center")
    ax.axvline(macro, color=RED, lw=1.5, ls=":", zorder=2)
    ax.text(macro - 0.7, -0.72, f"held-out macro avg {macro:.1f}%",
            color=RED, fontsize=9.5, ha="right", va="center")
    ax.set_xlim(0, 100)
    ax.set_ylim(-1.05, len(names) - 0.3 + 0.55)
    ax.set_xlabel("per-second F1 (%)")
    ax.set_title("Shipped config on the 5 films the optimizer never saw",
                 loc="left", fontsize=12, pad=12)
    fig.tight_layout()
    fig.savefig(out, dpi=160)
    plt.close(fig)


def fig_rep4_matrix(out: Path):
    combos = []
    for path in sorted(RESULTS.glob("rep4_best_*.json")):
        stem = path.stem[len("rep4_best_"):]
        for slug in MODEL_COLOURS:
            if stem.startswith(slug):
                mode = stem[len(slug) + 1:]  # e.g. full_exp
                with open(path) as f:
                    best = json.load(f)["best"]
                combos.append((slug, mode, best["f1"] * 100))
                break
    combos.sort(key=lambda c: c[2])

    fig, ax = plt.subplots(figsize=(9, 6.2))
    ax.grid(axis="y", visible=False)
    labels = []
    for i, (slug, mode, f1) in enumerate(combos):
        label, colour = MODEL_COLOURS[slug]
        scope, exp = mode.rsplit("_", 1)
        labels.append(f"{scope} · {'expand' if exp == 'exp' else 'no expand'}")
        ax.hlines(i, 50, f1, color=GRID, lw=1.2, zorder=2)
        ax.plot(f1, i, "o", ms=9, color=colour, zorder=3,
                mfc=colour if scope == "restricted" else SURFACE,
                mec=colour, mew=2)
        ax.text(f1 + 0.35, i, f"{f1:.1f}", va="center", fontsize=8.5, color=MUTED)
    ax.set_yticks(range(len(combos)), labels, fontsize=9)
    ax.set_xlim(65, 80)
    ax.set_xlabel("training-set per-second F1 (%)")
    ax.set_title("All 16 combos — filled dot = cast-restricted gallery, open = full",
                 loc="left", fontsize=12, pad=12)
    handles = [plt.Line2D([], [], marker="o", ls="", ms=9, color=c, label=l)
               for _, (l, c) in MODEL_COLOURS.items()]
    ax.legend(handles=handles, loc="lower right", frameon=False, fontsize=9.5)
    fig.tight_layout()
    fig.savefig(out, dpi=160)
    plt.close(fig)


def fig_de_landscape(out: Path):
    evals = []
    with open(REPO / "experiments/trajectories/rep4_LVFace-B_Glint360K_full_exp.jsonl") as f:
        for line in f:
            d = json.loads(line)
            evals.append((d["config"]["prob_threshold"],
                          d["config"]["extinction_sec"], d["f1"] * 100))
    x, y, f1 = map(np.array, zip(*evals))
    best = training_best()

    # one-hue sequential ramp (light -> dark blue), per the report palette
    cmap = LinearSegmentedColormap.from_list(
        "seq_blue", ["#cde2fb", "#86b6ef", "#3987e5", "#1c5cab", "#0d366b"])

    fig, ax = plt.subplots(figsize=(9, 5.2))
    # clip the color scale to the top of the range — DE spends most evals near
    # the optimum, so an unclipped scale renders the structure invisible
    sc = ax.scatter(x, y, c=f1, cmap=cmap, s=22, linewidths=0, zorder=3,
                    vmin=70, vmax=float(f1.max()))
    ax.plot(best["config"]["prob_threshold"], best["config"]["extinction_sec"],
            marker="*", ms=18, color=RED, mec=SURFACE, mew=1.2, zorder=4)
    ax.annotate(f"shipped optimum  F1 {best['f1']*100:.1f}%",
                (best["config"]["prob_threshold"], best["config"]["extinction_sec"]),
                textcoords="offset points", xytext=(-14, -30),
                ha="right", fontsize=10, color=RED,
                arrowprops={"arrowstyle": "-", "color": RED, "lw": 1})
    cb = fig.colorbar(sc, ax=ax, pad=0.02)
    cb.set_label("per-second F1 (%)")
    cb.outline.set_visible(False)
    ax.set_xlabel("prob_threshold")
    ax.set_ylabel("extinction_sec")
    ax.set_title("All 512 DE evaluations, LVFace-B full-gallery + expansion",
                 loc="left", fontsize=12, pad=12)
    fig.tight_layout()
    fig.savefig(out, dpi=160)
    plt.close(fig)


def fig_downton_timeline(out: Path, t0: int = 7100, t1: int = 7340):
    import h5py

    tracker = {}
    with open(RESULTS / "holdout/raw_Downton_Abbey__A_New_Era.jsonl") as f:
        for line in f:
            d = json.loads(line)
            tracker[int(d["timestamp_sec"])] = len(d["visible_actors"])
    with h5py.File(REPO / "experiments/dumps/LVFace-B_Glint360K/"
                          "dump_Downton_Abbey__A_New_Era.h5", "r") as h5:
        ts = h5["frames/timestamp_sec"][:]
        fc = h5["frames/face_count"][:]
    det = {int(t): int(c) for t, c in zip(ts, fc)}

    t = np.arange(t0, t1)
    # the raw stream occasionally skips a second under replay load — carry the
    # last seen value forward rather than dropping to 0
    trk, last = [], 0
    for s in t:
        if s in tracker:
            last = tracker[s]
        trk.append(last)
    trk = np.array(trk)
    dc = np.array([det.get(s, 0) for s in t])

    fig, ax = plt.subplots(figsize=(9.5, 4.4))
    ax.grid(axis="x", visible=False)
    ax.fill_between(t, dc, step="mid", color=GREEN, alpha=0.25, zorder=2)
    ax.step(t, dc, where="mid", color=GREEN, lw=2, zorder=3)
    ax.step(t, trk, where="mid", color=BLUE, lw=2, zorder=4)

    # longest contiguous run of "detector sees nothing, tracker still reporting"
    ghost = (dc == 0) & (trk > 0)
    runs, start = [], None
    for i, g in enumerate(ghost):
        if g and start is None:
            start = i
        elif not g and start is not None:
            runs.append((start, i - 1))
            start = None
    if start is not None:
        runs.append((start, len(ghost) - 1))
    if runs:
        i0, i1 = max(runs, key=lambda r: r[1] - r[0])
        g0, g1 = t[i0], t[i1]
        ax.axvspan(g0, g1, color=RED, alpha=0.08, zorder=1)
        ax.annotate(f"{g1 - g0}s of credits: 0 faces detected,\n"
                    f"{trk[i0]} actors still reported (frozen boxes)",
                    ((g0 + g1) / 2, 20.5), ha="center", va="bottom",
                    fontsize=10, color=RED)
    ax.text(t0 + 4, 27.3, "actors reported by tracker", color=BLUE,
            fontsize=10.5, va="bottom")
    ax.text(t0 + 4, 11.5, "faces seen by detector", color=GREEN,
            fontsize=10.5, va="bottom")
    ax.set_xlabel("film time (s)")
    ax.set_ylabel("count")
    ax.set_ylim(0, 31)
    ax.set_title("Downton Abbey: A New Era — the cut to credits, second by second",
                 loc="left", fontsize=12, pad=12)
    fig.tight_layout()
    fig.savefig(out, dpi=160)
    plt.close(fig)


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--out-dir", type=Path,
                   default=REPO / "docs/assets/images")
    args = p.parse_args()
    args.out_dir.mkdir(parents=True, exist_ok=True)

    fig_holdout_f1(args.out_dir / "holdout_f1_by_film.png")
    fig_rep4_matrix(args.out_dir / "rep4_matrix_f1.png")
    fig_de_landscape(args.out_dir / "de_search_landscape.png")
    fig_downton_timeline(args.out_dir / "downton_ghost_timeline.png")
    print(f"[experiment_charts] wrote 4 figures to {args.out_dir}")


if __name__ == "__main__":
    main()