"""Regenerate the Axis B signature figures from results/signature_benchmarks/bench_*.json.

    .venv/bin/python figures/make_sig_figures.py

Outputs (PDF, serif, publication quality):
  figures/sig_batch_verification.pdf  - per-signature time vs batch size for every MSM algorithm,
                                        one panel per host, single-verify reference lines (k256 and
                                        libsecp256k1); second row: speedup S(b) with bootstrap 95 % CI.
  figures/sig_implementations.pdf     - sign / verify ns per op across implementations (pure Python
                                        src/wp34, python-ecdsa, coincurve, k256, libsecp256k1), log scale.
Pass --copy-to-paper to also copy them into paper/figures/.
"""
from __future__ import annotations
import json, shutil
from pathlib import Path

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt  # noqa: E402
from matplotlib.ticker import FuncFormatter, LogLocator  # noqa: E402

ROOT = Path(__file__).resolve().parents[1]
RES = ROOT / "results" / "signature_benchmarks"
OUT = ROOT / "figures"

# fixed categorical order (validated with the dataviz palette checker: light surface, all pairs pass)
C = {"pippenger": "#2a78d6", "straus_wnaf5": "#eb6834", "k256_lincomb_vartime": "#1baf7a",
     "k256_lincomb_ct": "#6250d6", "single": "#52514e"}
ALG_LABEL = {"pippenger": "own Pippenger (bucket, mixed adds)", "straus_wnaf5": "own Straus (wNAF-5, interleaved)",
             "k256_lincomb_vartime": "k256 lincomb_vartime (GLV + wNAF)", "k256_lincomb_ct": "k256 lincomb (constant time)"}
ALG_ORDER = ["pippenger", "straus_wnaf5", "k256_lincomb_vartime", "k256_lincomb_ct"]

plt.rcParams.update({
    "font.family": "serif", "font.serif": ["Times New Roman", "Times", "DejaVu Serif"], "font.size": 9,
    "axes.titlesize": 9.5, "axes.labelsize": 9, "legend.fontsize": 7.5, "xtick.labelsize": 8, "ytick.labelsize": 8,
    "axes.spines.top": False, "axes.spines.right": False, "axes.grid": True, "grid.color": "#e4e3df",
    "grid.linewidth": 0.5, "axes.edgecolor": "#8a8984", "axes.linewidth": 0.6, "pdf.fonttype": 42,
    "figure.dpi": 150, "savefig.bbox": "tight",
})


def load():
    docs = [json.loads(f.read_text()) for f in sorted(RES.glob("bench_*.json"))]
    docs.sort(key=lambda d: d["label"])
    return docs


def short_label(d):
    return f"{d['label']}\n{d['host']['cpu']}, {d['host']['cores']} cores, macOS {d['host']['os'].split('(')[-1].rstrip(')')}"


def fig_batch(docs):
    n = len(docs)
    fig, axes = plt.subplots(2, n, figsize=(3.4 * n + 0.6, 5.6), sharex=True, sharey="row", squeeze=False)
    for j, d in enumerate(docs):
        m = d["measurements"]["batch_verification"]
        ax, ax2 = axes[0][j], axes[1][j]
        ref = m["single_verify_reference_ns"]
        ax.axhline(ref["libsecp256k1_schnorr_verify"] / 1e3, color=C["single"], lw=1.0, ls="--",
                   label="single verify, libsecp256k1 (reference)" if j == 0 else None)
        ax.axhline(ref["k256_verify_raw"] / 1e3, color=C["single"], lw=1.0, ls=":",
                   label="single verify, k256 verify_raw (reference)" if j == 0 else None)
        for alg in ALG_ORDER:
            rows = sorted([r for r in m["results"] if r["algorithm"] == alg], key=lambda r: r["batch_size"])
            if not rows:
                continue
            xs = [r["batch_size"] for r in rows]
            ys = [r["per_signature"]["median_ns"] / 1e3 for r in rows]
            lo = [r["per_signature"]["q1_ns"] / 1e3 for r in rows]
            hi = [r["per_signature"]["q3_ns"] / 1e3 for r in rows]
            ax.fill_between(xs, lo, hi, color=C[alg], alpha=0.18, lw=0)
            ax.plot(xs, ys, color=C[alg], lw=1.6, marker="o", ms=3.2, label=ALG_LABEL[alg])
            sp = [r["speedup_vs_k256_single_verify"] for r in rows]
            ax2.fill_between(xs, [s["ci95_low"] for s in sp], [s["ci95_high"] for s in sp], color=C[alg], alpha=0.18, lw=0)
            ax2.plot(xs, [s["point"] for s in sp], color=C[alg], lw=1.6, marker="o", ms=3.2)
            if alg == "pippenger":
                sp2 = [r["speedup_vs_libsecp256k1_single_verify"] for r in rows]
                ax2.plot(xs, [s["point"] for s in sp2], color=C[alg], lw=1.2, ls="--", marker="s", ms=2.8,
                         label="own Pippenger vs libsecp256k1 single verify")
                last = rows[-1]
                ax.annotate(f"{last['per_signature']['median_ns']/1e3:.1f} µs/sig", (xs[-1], ys[-1]), xytext=(-6, 6),
                            textcoords="offset points", ha="right", va="bottom", fontsize=7, color=C[alg])
        ax2.axhline(1.0, color=C["single"], lw=0.8, ls="-")
        ax.set_xscale("log", base=2)
        ax.set_yscale("log")
        ax.set_ylim(5, 160)
        ax.set_title(short_label(d), loc="left")
        ax.yaxis.set_major_formatter(FuncFormatter(lambda v, _: f"{v:g}"))
        ax.yaxis.set_minor_locator(LogLocator(base=10, subs=(2, 5)))
        ax.yaxis.set_minor_formatter(FuncFormatter(lambda v, _: f"{v:g}"))
        ax2.set_xlabel("batch size $b$ (signatures per batch)")
        ax2.set_xticks([1, 4, 16, 64, 256, 1024, 2048])
        ax2.set_xticklabels(["1", "4", "16", "64", "256", "1024", "2048"])
        ax2.set_ylim(0, 3.6)
        if j == 0:
            ax.set_ylabel("time per signature (µs, median; band = IQR)")
            ax2.set_ylabel("speed-up $S(b)$ vs single k256 verify\n(band = bootstrap 95 % CI)")
        else:
            ax.tick_params(labelleft=False)
            ax2.tick_params(labelleft=False)
    h1, l1 = axes[0][0].get_legend_handles_labels()
    h2, l2 = axes[1][0].get_legend_handles_labels()
    fig.legend(h1 + h2, l1 + l2, loc="lower center", ncol=3, frameon=False, bbox_to_anchor=(0.5, -0.06))
    fig.suptitle("BIP340 batch verification on secp256k1 (k256 group arithmetic; 128-bit random weights, $a_1 = 1$; synthetic keys, seed 34)",
                 fontsize=9.5, y=1.0)
    fig.tight_layout(h_pad=0.6)
    fig.savefig(OUT / "sig_batch_verification.pdf")
    plt.close(fig)


def fig_impl(docs):
    impls = [  # (label, ecdsa_sign, ecdsa_verify, schnorr_sign, schnorr_verify) -> keys in sign_verify or python baselines
        ("pure Python\nsrc/wp34", "py:wp34_pure_python_ecdsa_sign", "py:wp34_pure_python_ecdsa_verify",
         "py:wp34_pure_python_schnorr_sign", "py:wp34_pure_python_schnorr_verify"),
        ("python-ecdsa\n(pure Python)", "py:python_ecdsa_pkg_ecdsa_sign", "py:python_ecdsa_pkg_ecdsa_verify", None, None),
        ("coincurve\n(CFFI, libsecp256k1)", "py:coincurve_ecdsa_sign_der", "py:coincurve_ecdsa_verify_der",
         "py:coincurve_schnorr_sign", "py:coincurve_schnorr_verify"),
        ("k256 0.14\n(Rust)", "rs:ecdsa_sign_k256", "rs:ecdsa_verify_k256", "rs:schnorr_sign_k256", "rs:schnorr_verify_k256"),
        ("libsecp256k1\n(C, via Rust)", "rs:ecdsa_sign_libsecp256k1", "rs:ecdsa_verify_libsecp256k1",
         "rs:schnorr_sign_libsecp256k1", "rs:schnorr_verify_libsecp256k1"),
    ]
    series = [("ECDSA sign", 1, "#2a78d6", ""), ("ECDSA verify", 2, "#2a78d6", "////"),
              ("Schnorr sign", 3, "#eb6834", ""), ("Schnorr verify", 4, "#eb6834", "////")]
    n = len(docs)
    fig, axes = plt.subplots(1, n, figsize=(4.6 * n + 0.4, 3.7), sharey=True, squeeze=False)
    for j, d in enumerate(docs):
        ax = axes[0][j]
        sv = d["measurements"]["sign_verify"]
        py = (d["measurements"].get("python_baselines") or {}).get("results") or {}

        def val(key):
            if key is None:
                return None
            src, k = key.split(":", 1)
            t = sv.get(k) if src == "rs" else py.get(k)
            return None if t is None else t["median_ns"]

        w = 0.2
        for s_i, (sname, col, color, hatch) in enumerate(series):
            xs, ys = [], []
            for i, imp in enumerate(impls):
                v = val(imp[col])
                if v is not None:
                    xs.append(i + (s_i - 1.5) * w)
                    ys.append(v)
            bars = ax.bar(xs, ys, width=w * 0.92, color=color, hatch=hatch, edgecolor="white", linewidth=0.6,
                          label=sname if j == 0 else None, alpha=0.95 if not hatch else 0.75)
            for b_, v in zip(bars, ys):
                txt = f"{v/1e3:.0f}" if v >= 1e5 else f"{v/1e3:.1f}"
                ax.text(b_.get_x() + b_.get_width() / 2, v * 1.12, txt, ha="center", va="bottom", fontsize=5.8, rotation=90, color="#0b0b0b")
        for i, imp in enumerate(impls):
            if all(val(k) is None for k in imp[1:]):
                ax.text(i, 2.5e4, "not installed\non this host", ha="center", va="center", fontsize=6.5, color="#52514e")
        ax.set_yscale("log")
        ax.set_ylim(3e3, 5e7)
        ax.set_xticks(range(len(impls)))
        ax.set_xticklabels([i[0] for i in impls], fontsize=6.8)
        ax.set_title(short_label(d), loc="left")
        ax.yaxis.set_major_formatter(FuncFormatter(lambda v, _: f"{v/1e3:g}"))
        if j == 0:
            ax.set_ylabel("time per operation (µs, median of ≥ 7 reps; log scale)")
        ax.grid(axis="x", visible=False)
    fig.legend(loc="lower center", ncol=4, frameon=False, bbox_to_anchor=(0.5, -0.06))
    fig.suptitle("Implementation spread for secp256k1 signatures (32-byte prehash; Rust: 2000 ops/rep, Python: 20 ops/rep; labels in µs)",
                 fontsize=9.5, y=1.0)
    fig.tight_layout()
    fig.savefig(OUT / "sig_implementations.pdf")
    plt.close(fig)


def main():
    import sys
    docs = load()
    if not docs:
        raise SystemExit("no results/signature_benchmarks/bench_*.json")
    fig_batch(docs)
    fig_impl(docs)
    paper = ROOT / "paper" / "figures"
    if "--copy-to-paper" in sys.argv and paper.exists():
        for f in ("sig_batch_verification.pdf", "sig_implementations.pdf"):
            shutil.copy(OUT / f, paper / f)
    print("wrote", OUT / "sig_batch_verification.pdf", OUT / "sig_implementations.pdf")


if __name__ == "__main__":
    main()