"""Generate figures from results/*.json into paper/figures/."""
from __future__ import annotations

import json
from pathlib import Path

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

ROOT = Path(__file__).resolve().parent.parent
RESULTS = ROOT / "results"
FIGS = ROOT / "figures"

ARM_STYLE = {
    "no_memory": ("No memory (schema draft)", "#888888", "o", "--"),
    "static_global": ("Static datastore (ToolSpec-style)", "#d1495b", "s", "-"),
    "personal_memory": ("Ours (personal + evict)", "#1b6ca8", "D", "-"),
}


def plot_acceptance(tag=""):
    fname = f"{tag}_accept_results.json" if tag else "accept_results.json"
    data = json.loads((RESULTS / fname).read_text())
    summary = data["summary"]
    n_sessions = data["config"]["sessions"]

    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(11, 4.2))
    for arm, (label, color, marker, ls) in ARM_STYLE.items():
        if arm not in summary:
            continue
        xs = sorted(int(s) for s in summary[arm])
        mat = [summary[arm][str(s)]["MAT"] for s in xs]
        frac = [summary[arm][str(s)]["accepted_frac"] for s in xs]
        ax1.plot(xs, mat, marker=marker, ls=ls, color=color, label=label)
        ax2.plot(xs, frac, marker=marker, ls=ls, color=color, label=label)

    for ax, ylab, title in ((ax1, "Mean accepted tokens (MAT)",
                             "Draft acceptance vs. session"),
                            (ax2, "Accepted fraction of target",
                             "Accepted fraction vs. session")):
        ax.set_xlabel("Session index")
        ax.set_ylabel(ylab)
        ax.set_title(title)
        ax.grid(alpha=0.3)
        ax.axvspan(-0.4, 0.4, color="k", alpha=0.05)
    ax1.legend(fontsize=8, loc="best")
    ax1.annotate("warmup", (0, ax1.get_ylim()[1]*0.05), fontsize=7, ha="center")
    fig.tight_layout()
    FIGS.mkdir(parents=True, exist_ok=True)
    fig.savefig(FIGS / "acceptance_by_session.pdf")
    fig.savefig(FIGS / "acceptance_by_session.png", dpi=140)
    print("wrote acceptance_by_session.{pdf,png}")


def plot_safety():
    """Two-panel safety figure over the three execution policies:
    (left) severity-weighted cost of being wrong; (right) safe spec-executions
    preserved (the latency win) and bad irreversible actions."""
    data = json.loads((RESULTS / "safety_results.json").read_text())
    policies = ["naive_exec", "conf_gate", "gated_exec"]
    labels = ["Naive\nspec-exec", "Confidence\ngate", "Idempotency\ngate (ours)"]
    colors = ["#d1495b", "#e8a33d", "#1b6ca8"]

    wcost = [data[p]["weighted_cost"] for p in policies]
    bad = [data[p]["bad_irreversible_actions"] for p in policies]
    safe = [data[p]["safe_spec_executions"] for p in policies]
    x = list(range(len(policies)))

    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 4.1))

    ax1.bar(x, wcost, color=colors)
    for i, w in enumerate(wcost):
        ax1.text(i, w + max(wcost) * 0.02, f"{w:.0f}", ha="center",
                 fontsize=10, fontweight="bold", color=colors[i])
    ax1.set_xticks(x)
    ax1.set_xticklabels(labels, fontsize=8.5)
    ax1.set_ylabel("Severity-weighted cost of being wrong")
    ax1.set_title("Cost of wrong irreversible speculative executions")
    ax1.grid(axis="y", alpha=0.3)

    ax2.bar(x, safe, color="#1b6ca8", label="Safe spec-executions (latency win)")
    ax2.bar(x, bad, bottom=safe, color="#d1495b",
            label="Bad irreversible actions")
    for i, b in enumerate(bad):
        ax2.text(i, safe[i] + b + max(safe) * 0.02,
                 f"{b:.0f} bad", ha="center", fontsize=8.5,
                 color="#d1495b", fontweight="bold")
    ax2.set_xticks(x)
    ax2.set_xticklabels(labels, fontsize=8.5)
    ax2.set_ylabel("Speculative executions (mean/stream)")
    ax2.set_title("Latency win preserved vs. harm incurred")
    ax2.legend(fontsize=8, loc="upper right")
    ax2.grid(axis="y", alpha=0.3)

    fig.tight_layout()
    FIGS.mkdir(parents=True, exist_ok=True)
    fig.savefig(FIGS / "safety.pdf")
    fig.savefig(FIGS / "safety.png", dpi=140)
    print("wrote safety.{pdf,png}")


if __name__ == "__main__":
    import sys
    tag = sys.argv[1] if len(sys.argv) > 1 else ""
    plot_acceptance(tag)
    plot_safety()