"""Figures for the how-it-works page, from the demo data.

    python tools/make_figures.py        # writes out/*.png and prints the numbers quoted on the page
"""
from __future__ import annotations

import json
import os
import sys

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, ROOT)
from router import bounds  # noqa: E402
from router.data import load_csv  # noqa: E402
from router.plan import plan  # noqa: E402

OUT = os.path.join(ROOT, "out")
COLORS = ["#eb6834", "#1baf7a", "#eda100", "#e87ba4", "#6250d6"]
ACCENT, TEAL, INK = "#2153D0", "#0E8C7A", "#0E1A2E"


def frontier(rep):
    fig, ax = plt.subplots(figsize=(8, 5))
    k = 1000
    orc = rep["oracle"]["hold"]["curve"]
    ax.plot([p["cost"] * k for p in orc], [100 * p["quality"] for p in orc], "--", color=INK, lw=1.3, label="oracle: cheapest model that is right")
    sens = rep["sensitivity"]
    cs = [r for r in sens if r.get("cascade")]
    ax.plot([r["cascade"]["cost"] * k for r in cs], [100 * r["cascade"]["quality"] for r in cs], "-", color=ACCENT, lw=1.6, label="cascade, one point per target")
    for r in cs:
        ax.plot(r["cascade"]["cost"] * k, 100 * r["cascade"]["quality"], "o", ms=4.5, mfc=ACCENT if r["cascade"]["certified"] else "white", mec=ACCENT)
    rs = [r for r in sens if r.get("router")]
    ax.plot([r["router"]["cost"] * k for r in rs], [100 * r["router"]["quality"] for r in rs], "-", color=TEAL, lw=1.4, label="learned router (prompt text)")
    for i, s in enumerate(rep["singles"]):
        h = s["hold"]
        ax.plot(h["cost"] * k, 100 * h["quality"], "o", ms=9, color=COLORS[i % len(COLORS)], zorder=5)
        ax.annotate(s["name"] + " alone", (h["cost"] * k, 100 * h["quality"]), textcoords="offset points", xytext=(9, -3), fontsize=10)
    ax.axhline(100 * rep["target"], color=ACCENT, ls=":", lw=1)
    c = rep["cascade"]["hold"]
    ax.plot(c["cost"] * k, 100 * c["quality"], "o", ms=14, mfc="none", mec=ACCENT, mew=2)
    ax.set_xscale("log")
    ax.set_xlabel("cost per 1,000 prompts ($)")
    ax.set_ylabel("quality on the 750 held-out prompts (%)")
    ax.set_title("Cost against quality: filled cascade points are certified at their target", fontsize=11)
    ax.grid(alpha=.3, which="both")
    ax.legend(loc="lower right", fontsize=9)
    fig.tight_layout()
    fig.savefig(os.path.join(OUT, "frontier.png"), dpi=130)
    plt.close(fig)


def sensitivity(rep):
    fig, ax = plt.subplots(figsize=(8, 4.6))
    k = 1000
    sens = rep["sensitivity"]
    t = [100 * r["target"] for r in sens]
    ax.plot([x for x, r in zip(t, sens) if r.get("oracleCost") is not None], [r["oracleCost"] * k for r in sens if r.get("oracleCost") is not None], "--", color=INK, lw=1.3, label="oracle")
    bs = [(x, r["bestSingle"]["cost"] * k) for x, r in zip(t, sens) if r.get("bestSingle")]
    ax.step([b[0] for b in bs], [b[1] for b in bs], where="post", color="#eda100", lw=1.6, label="cheapest certified single model")
    ax.plot(t, [r["router"]["cost"] * k for r in sens], "-", color=TEAL, lw=1.4, label="learned router")
    ax.plot(t, [r["cascade"]["cost"] * k for r in sens], "-", color=ACCENT, lw=1.6, label="cascade")
    for x, r in zip(t, sens):
        ax.plot(x, r["cascade"]["cost"] * k, "o", ms=4.5, mfc=ACCENT if r["cascade"]["certified"] else "white", mec=ACCENT)
    ax.axvline(100 * rep["target"], color=ACCENT, ls=":", lw=1)
    ax.set_yscale("log")
    ax.set_xlabel("quality target (%)")
    ax.set_ylabel("held-out cost per 1,000 prompts ($)")
    ax.set_title("The price of each extra point of quality (hollow = not certified at that target)", fontsize=11)
    ax.grid(alpha=.3, which="both")
    ax.legend(fontsize=9, loc="upper left")
    fig.tight_layout()
    fig.savefig(os.path.join(OUT, "sensitivity.png"), dpi=130)
    plt.close(fig)


def overfit(rep, log):
    """Calibration quality the search hit vs. what the same thresholds scored held out."""
    from router.policies import evaluate_cascade, split_indices
    cal, hold = split_indices(log.n, rep["calFrac"], rep["seed"])
    fig, ax = plt.subplots(figsize=(6.4, 4.6))
    xs, ys, aims = [], [], []
    for r in rep["sensitivity"]:
        c = evaluate_cascade(log, cal, r["thresholds"], log.confidence)
        xs.append(100 * c["quality"]); ys.append(100 * r["cascade"]["quality"]); aims.append(100 * (r["aim"] or 1))
    ax.plot([60, 100], [60, 100], color="#999", lw=1, ls="--", label="no drop")
    ax.scatter(xs, ys, color=ACCENT, s=28, label="one plan per target")
    ax.set_xlabel("quality on the calibration prompts (what the search saw)")
    ax.set_ylabel("quality on the held-out prompts")
    ax.set_title("Tuned thresholds look better on the rows they were tuned on", fontsize=11)
    ax.grid(alpha=.3)
    ax.legend(fontsize=9)
    fig.tight_layout()
    fig.savefig(os.path.join(OUT, "overfit.png"), dpi=130)
    plt.close(fig)
    d = np.array(xs) - np.array(ys)
    return float(d.mean()), float(d.max())


def bound_width():
    fig, axes = plt.subplots(1, 2, figsize=(9, 4))
    ns = np.array([50, 100, 200, 400, 750, 1500, 3000, 6000, 12000])
    for p, col in [(0.8, COLORS[0]), (0.9, ACCENT), (0.95, TEAL)]:
        gap = [100 * (p - bounds.clopper_pearson_lower(round(p * n), n, 0.05)) for n in ns]
        axes[0].plot(ns, gap, "o-", ms=4, color=col, label=f"observed {int(100 * p)}%")
    axes[0].set_xscale("log"); axes[0].set_yscale("log")
    axes[0].set_xlabel("held-out prompts"); axes[0].set_ylabel("observed rate minus 95% lower bound (points)")
    axes[0].set_title("How much the bound gives away", fontsize=11); axes[0].grid(alpha=.3, which="both"); axes[0].legend(fontsize=9)
    rng = np.random.default_rng(0)
    n, sims = 200, 4000
    for p, col in [(0.8, COLORS[0]), (0.9, ACCENT), (0.95, TEAL)]:
        ks = rng.binomial(n, p, size=sims)
        cache = {k: bounds.clopper_pearson_lower(int(k), n, 0.05) for k in np.unique(ks)}
        lbs = np.array([cache[k] for k in ks])
        axes[1].hist(100 * lbs, bins=40, alpha=.55, color=col, label=f"true {int(100 * p)}%: bound above truth {100 * (lbs > p).mean():.1f}% of the time")
        axes[1].axvline(100 * p, color=col, lw=1)
    axes[1].set_xlabel("95% lower bound from 200 held-out prompts (%)"); axes[1].set_ylabel("simulations")
    axes[1].set_title("The bound fails at most 5% of the time", fontsize=11); axes[1].legend(fontsize=8)
    fig.tight_layout()
    fig.savefig(os.path.join(OUT, "bounds.png"), dpi=130)
    plt.close(fig)


def main():
    os.makedirs(OUT, exist_ok=True)
    log = load_csv(os.path.join(ROOT, "web", "demo.csv"))
    rep = plan(log, target=0.9, delta=0.05)
    frontier(rep)
    sensitivity(rep)
    mean_drop, max_drop = overfit(rep, log)
    bound_width()
    with open(os.path.join(OUT, "demo_plan.json"), "w", encoding="utf-8") as f:
        json.dump(rep, f, indent=1)
    c, r, o = rep["cascade"]["hold"], rep["router"]["hold"], rep["oracle"]["hold"]
    fr = rep["singles"][-1]["hold"]
    print(f"aim {rep['aim']:.4f} (required rate {rep['requiredRate']:.4f}); enumerated {rep['enumerated']} cascades on a grid of {rep['gridSize']}")
    print("thresholds", rep["cascade"]["thresholds"])
    print(f"cascade: quality {c['quality']:.4f} LB {c['qualityLB']:.4f} cost/1k {1000 * c['cost']:.3f} ({c['cost'] / fr['cost']:.3f} of frontier) UB {1000 * c['costUB']:.3f} latency {c['latency']:.0f}")
    print("  shares", [round(x, 3) for x in c["shares"]], "calls", [round(x, 3) for x in c["calls"]])
    print(f"router: quality {r['quality']:.4f} LB {r['qualityLB']:.4f} cost/1k {1000 * r['cost']:.3f} tau {rep['router']['tau']:.3f}")
    print(f"oracle: quality {o['quality']:.4f} cost/1k {1000 * o['cost']:.3f} shares {[round(x, 3) for x in o['shares']]} none right {o['noneRight']:.4f}")
    for s in rep["singles"]:
        h = s["hold"]
        print(f"{s['name']:9s} quality {h['quality']:.4f} LB {h['qualityLB']:.4f} cost/1k {1000 * h['cost']:.3f} latency {h['latency']:.0f} certified {h['certified']}")
    print(f"overfit drop cal->hold: mean {mean_drop:.2f} points, max {max_drop:.2f}")
    print("breakEven", rep["breakEven"])
    cert = [(r['target'], round(1000 * r['cascade']['cost'], 3), r['cascade']['certified']) for r in rep["sensitivity"]]
    print("sweep", cert)


if __name__ == "__main__":
    main()