#!/usr/bin/env python3
"""Generate every figure of the paper from logged measurements (benchmarks/raw, benchmarks/processed). PDF outputs in
paper/figures/ and figures/. Nothing here is hand-drawn from numbers typed in."""
import glob, json, os, statistics, sys
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
FIG = os.path.join(ROOT, "paper", "figures"); FIG2 = os.path.join(ROOT, "figures")
os.makedirs(FIG, exist_ok=True); os.makedirs(FIG2, exist_ok=True)
plt.rcParams.update({"font.size": 9, "axes.spines.top": False, "axes.spines.right": False, "figure.dpi": 150})
PAL = ["#2E5A88", "#C8553D", "#4F8A5B", "#8E6C8A", "#D9A441", "#5B5B5B", "#7FB2D6"]

def save(fig, name):
    for d in (FIG, FIG2):
        fig.savefig(os.path.join(d, name + ".pdf"), bbox_inches="tight")
    plt.close(fig); print("wrote", name)

def fig_sim_hit_vs_capacity():
    p = os.path.join(ROOT, "benchmarks/processed/gateC/sim_olmoe_v2_ssd043_m8.json")
    if not os.path.exists(p): return
    d = json.load(open(p)); rows = [r for r in d["results"] if r["split"] == "heldout" and r["policy"] == "lru" and r["predictor"] != "probe"]
    for extra in ("sim_olmoe_v2_ssd043_m8_probe_fixed.json", "sim_olmoe_v2_ssd043_m8_probe_depths.json"):
        q = os.path.join(ROOT, "benchmarks/processed/gateC", extra)
        if os.path.exists(q):
            rows += [r for r in json.load(open(q))["results"] if r["split"] == "heldout" and r["policy"] == "lru" and r["predictor"] not in {x["predictor"] for x in rows}]
    preds = []
    for r in rows:
        if r["predictor"] not in preds: preds.append(r["predictor"])
    fig, axes = plt.subplots(1, 2, figsize=(7.2, 2.8))
    for i, pr in enumerate(preds):
        rs = sorted([r for r in rows if r["predictor"] == pr], key=lambda r: r["capacity"])
        axes[0].plot([r["capacity"] for r in rs], [r["hit_rate"] for r in rs], marker="o", ms=3, color=PAL[i % len(PAL)], label=pr)
        axes[1].plot([r["capacity"] for r in rs], [r["stall_ms_per_token"] for r in rs], marker="o", ms=3, color=PAL[i % len(PAL)], label=pr)
    axes[0].set_xlabel("slots per layer $C$"); axes[0].set_ylabel("expert hit rate (held-out)"); axes[0].set_ylim(0.4, 1.0)
    axes[1].set_xlabel("slots per layer $C$"); axes[1].set_ylabel("modelled stall (ms / token)")
    axes[1].legend(frameon=False, fontsize=7, ncol=2)
    fig.suptitle(f"OLMoE-1B-7B, trace-driven simulation (t_read {d['t_read_ms']} ms, {d['loaders']} loader)", fontsize=9)
    save(fig, "sim_hit_vs_capacity")

def fig_capacity_sweep():
    paths = [os.path.join(ROOT, "benchmarks/raw/gateB/capacity_sweep_laptop/sweep.jsonl"),
             os.path.join(ROOT, "benchmarks/raw/gateC/prefetch_online_laptop/smoke_lazyinstall.jsonl"),
             os.path.join(ROOT, "benchmarks/raw/gateC/cold_nvme_laptop/sweep.jsonl")]
    recs = []
    for p in paths:
        if os.path.exists(p):
            recs += [dict(json.loads(l), src=os.path.basename(os.path.dirname(p))) for l in open(p) if l.strip()]
    recs = [r for r in recs if not r.get("warmup")]
    if not recs: return
    groups = {}
    for r in recs:
        key = (r["src"], r["backend"], r.get("capacity"), r.get("predictor") or "none", r.get("m"), bool(r.get("nocache")), r.get("tag", ""))
        groups.setdefault(key, []).append(r["gen_tps"])
    labels, vals, errs, cols = [], [], [], []
    for k, v in sorted(groups.items(), key=lambda kv: (kv[0][5], -(kv[0][2] or 999), kv[0][3])):
        src, be, cap, pr, m, cold, tag = k
        if be == "mlx_lm": lab = "mlx-lm (all resident)"
        else: lab = f"C={cap} {pr}{' m='+str(m) if pr!='none' else ''}{' cold' if cold else ' page-cache'} [{src.split('_')[0]}{(' '+tag) if tag else ''}]"
        labels.append(lab); vals.append(statistics.median(v)); errs.append((max(v) - min(v)) / 2); cols.append(PAL[1] if cold else PAL[0])
    fig, ax = plt.subplots(figsize=(7.2, 0.22 * len(labels) + 1.2))
    y = range(len(labels)); ax.barh(y, vals, xerr=errs, color=cols, height=0.7, error_kw={"lw": 0.8})
    ax.set_yticks(list(y)); ax.set_yticklabels(labels, fontsize=7); ax.invert_yaxis(); ax.set_xlabel("decode tok/s (median; bar = range)")
    ax.set_title("OLMoE-1B-7B 4-bit on M5 Max: residency, tier and prefetch", fontsize=9)
    save(fig, "capacity_sweep")

def fig_two_node():
    p = os.path.join(ROOT, "benchmarks/raw/gateE/two_node_m3u96a_m3u96b.jsonl")
    if not os.path.exists(p): return
    recs = [json.loads(l) for l in open(p) if l.strip()]; recs = [r for r in recs if not r["warmup"]]
    groups = {}
    for r in recs: groups.setdefault((r["capacity"], r["placement"]), []).append(r)
    fig, ax = plt.subplots(figsize=(4.2, 2.6))
    labels = []; i = 0
    for (cap, pl), v in sorted(groups.items()):
        tps = [x["gen_tps"] for x in v]; rw = [x["remote"]["remote_wait_s"] / max(1, x["gen_tokens"]) * 1e3 if x.get("remote") else 0 for x in v]
        ax.bar(i, statistics.median(tps), color=PAL[0 if pl == "none" else 1]); labels.append(f"C={cap}\n{pl}"); i += 1
    ax.set_xticks(range(i)); ax.set_xticklabels(labels, fontsize=7); ax.set_ylabel("decode tok/s (median)")
    ax.set_title("OLMoE: M3U96a local vs odd experts on M3U96b (1 GbE)", fontsize=8)
    save(fig, "two_node")

def fig_attribution():
    p = os.path.join(ROOT, "benchmarks/raw/gateC/cold_nvme_laptop/attribution.jsonl")
    if not os.path.exists(p): return
    recs = [json.loads(l) for l in open(p) if l.strip()]; recs = [r for r in recs if not r["warmup"]]
    groups = {}
    for r in recs: groups.setdefault((r["predictor"], r["m"]), []).append(r)
    order = sorted(groups, key=lambda k: (k[0] != "none", k[0], k[1]))
    comps = ["sync", "prefetch", "ensure"]; names = {"sync": "GPU sync (graph incl. installs)", "prefetch": "prefetch bookkeeping + late waits", "ensure": "demand reads (SSD)"}
    fig, ax = plt.subplots(figsize=(7.2, 3.0)); bottom = [0.0] * len(order)
    for ci, c in enumerate(comps):
        vals = []
        for k in order:
            v = groups[k]; n = sum(x["gen_tokens"] for x in v); vals.append(sum(x["host_time_s"][c] for x in v) / n * 1e3)
        ax.bar(range(len(order)), vals, bottom=bottom, color=PAL[ci], label=names[c]); bottom = [b + v for b, v in zip(bottom, vals)]
    for i, k in enumerate(order):
        v = groups[k]; ax.text(i, bottom[i] + 0.5, f"{statistics.median(x['gen_tps'] for x in v):.0f} tok/s\nhit {statistics.median(x['cache_hit_rate'] for x in v):.2f}", ha="center", fontsize=6)
    ax.set_xticks(range(len(order))); ax.set_xticklabels([f"{k[0]}\nm={k[1]}" for k in order], fontsize=6.5); ax.set_ylabel("host time per token (ms)")
    ax.set_ylim(0, max(bottom) * 1.25); ax.legend(frameon=False, fontsize=7, loc="upper left")
    ax.set_title("OLMoE, 16 slots, cold NVMe (M5 Max): where the decode step waits", fontsize=9)
    save(fig, "attribution")

if __name__ == "__main__":
    fig_sim_hit_vs_capacity(); fig_capacity_sweep(); fig_two_node(); fig_attribution()