#!/usr/bin/env python3
"""Break-even analysis of bench_primitives.jsonl: medians, bootstrap 95 % CIs, speedups T_ref/T_metal, break-even
elements per kernel, CPU-wins table; writes experiments/results/bench_primitives_summary.{json,csv} and figures.
Usage: analyze_bench_primitives.py [raw.jsonl] [out_dir]
"""
import json, os, sys, collections
import numpy as np

raw = sys.argv[1] if len(sys.argv) > 1 else "experiments/results/raw/bench_primitives.jsonl"
out_dir = sys.argv[2] if len(sys.argv) > 2 else "experiments/results"
recs = [json.loads(l) for l in open(raw) if l.strip()]
lib = [r for r in recs if r.get("kernel") == "shader_library_load"]
recs = [r for r in recs if r.get("kernel") != "shader_library_load"]
rng = np.random.default_rng(20261102)

def boot_ci(x, n=2000):
    x = np.asarray(x); meds = [np.median(rng.choice(x, x.size, replace=True)) for _ in range(n)]
    return float(np.median(x)), float(np.percentile(meds, 2.5)), float(np.percentile(meds, 97.5))

groups = collections.defaultdict(list)
for r in recs:
    key = (r["kernel"], r["S"], r["n"], r["backend"], r["threads"])
    groups[key].append(r)
rows = []
for (kernel, S, n, backend, threads), rs in sorted(groups.items()):
    w = [r["wall_seconds"] for r in rs]; med, lo, hi = boot_ci(w)
    dev = float(np.median([r["device_seconds"] for r in rs])) if backend == "metal" else 0.0
    rows.append({"kernel": kernel, "S": S, "n": n, "elements": S * n, "backend": backend, "threads": threads, "median_s": med, "ci_lo": lo, "ci_hi": hi,
                 "device_s": dev, "elements_per_s": S * n / med, "reps": len(rs), "precision": rs[0]["precision"], "peak_rss_mb": max(r["peak_rss_after"] for r in rs) / 1e6})
# speedups vs CPU multi-thread (reference) and CPU single-thread
by = {(r["kernel"], r["S"], r["n"], r["backend"], r["threads"] == 1): r for r in rows}
summary = []
for r in rows:
    if r["backend"] != "metal": continue
    cpuN = by.get((r["kernel"], r["S"], r["n"], "cpu", False)); cpu1 = by.get((r["kernel"], r["S"], r["n"], "cpu", True))
    summary.append({**r, "speedup_vs_cpu_mt": cpuN["median_s"] / r["median_s"] if cpuN else None, "speedup_vs_cpu_1t": cpu1["median_s"] / r["median_s"] if cpu1 else None,
                    "cpu_mt_median_s": cpuN["median_s"] if cpuN else None, "cpu_1t_median_s": cpu1["median_s"] if cpu1 else None,
                    "metal_wins_vs_mt": bool(cpuN and r["ci_hi"] < cpuN["ci_lo"]), "cpu_wins_vs_mt": bool(cpuN and cpuN["ci_hi"] < r["ci_lo"])})
# break-even: smallest element count where Metal beats CPU-MT (CI-separated) per kernel
breakeven = {}
for k in sorted(set(r["kernel"] for r in summary)):
    wins = sorted([r["elements"] for r in summary if r["kernel"] == k and r["metal_wins_vs_mt"]])
    losses = sorted([r["elements"] for r in summary if r["kernel"] == k and r["cpu_wins_vs_mt"]])
    breakeven[k] = {"smallest_metal_win_elements": wins[0] if wins else None, "largest_cpu_win_elements": losses[-1] if losses else None,
                    "max_speedup_vs_mt": max((r["speedup_vs_cpu_mt"] or 0) for r in summary if r["kernel"] == k), "n_cases": sum(1 for r in summary if r["kernel"] == k)}
os.makedirs(out_dir, exist_ok=True)
json.dump({"source": raw, "rows": rows, "metal_summary": summary, "breakeven": breakeven, "shader_library": lib}, open(os.path.join(out_dir, "bench_primitives_summary.json"), "w"), indent=1)
import csv
with open(os.path.join(out_dir, "bench_primitives_summary.csv"), "w", newline="") as f:
    w = csv.DictWriter(f, fieldnames=list(rows[0].keys())); w.writeheader(); w.writerows(rows)
print("kernel                 S      n     elements   cpu1(ms)  cpuMT(ms)  metal(ms)  dev(ms)  speedup_MT  verdict")
for r in sorted(summary, key=lambda r: (r["kernel"], r["elements"], r["S"])):
    v = "METAL" if r["metal_wins_vs_mt"] else ("CPU" if r["cpu_wins_vs_mt"] else "tie")
    print(f"{r['kernel']:20s} {r['S']:6d} {r['n']:7d} {r['elements']:11d} {1e3*(r['cpu_1t_median_s'] or 0):9.3f} {1e3*(r['cpu_mt_median_s'] or 0):9.3f} {1e3*r['median_s']:9.3f} {1e3*r['device_s']:8.3f} {r['speedup_vs_cpu_mt'] or 0:10.2f}  {v}")
print("break-even:", json.dumps(breakeven, indent=1))
if lib: print("shader library:", lib)
try:
    import matplotlib; matplotlib.use("Agg"); import matplotlib.pyplot as plt
    kernels = sorted(set(r["kernel"] for r in rows)); fig, axes = plt.subplots(1, len(kernels), figsize=(4.2 * len(kernels), 3.6), squeeze=False)
    for ax, k in zip(axes[0], kernels):
        for backend, thr, lab, col in [("cpu", True, "CPU 1 thread", "#9aa0a6"), ("cpu", False, "CPU multicore", "#1f4e79"), ("metal", False, "Metal (FP32 df)", "#c0392b")]:
            pts = sorted([(r["elements"], r["elements_per_s"], r["ci_lo"], r["ci_hi"], r["median_s"]) for r in rows if r["kernel"] == k and r["backend"] == backend and (r["threads"] == 1) == thr])
            if not pts: continue
            x = [p[0] for p in pts]; y = [p[1] for p in pts]
            ax.plot(x, y, "o-", ms=3, lw=1.2, label=lab, color=col)
        ax.set_xscale("log"); ax.set_yscale("log"); ax.set_title(k, fontsize=9); ax.set_xlabel("elements (S×n)"); ax.grid(alpha=0.3)
    axes[0][0].set_ylabel("elements / s (median)"); axes[0][0].legend(fontsize=7)
    fig.tight_layout(); fig.savefig(os.path.join(out_dir, "fig_bench_primitives_throughput.pdf")); fig.savefig(os.path.join(out_dir, "fig_bench_primitives_throughput.png"), dpi=150)
    print("figure written")
except Exception as e:
    print("figure skipped:", e)