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

OUT = "/root/alblit/diffusion/figures"
os.makedirs(OUT, exist_ok=True)
plt.rcParams.update({"font.family": "DejaVu Sans", "font.size": 9,
                     "axes.spines.top": False, "axes.spines.right": False,
                     "axes.edgecolor": "#444", "figure.dpi": 200})
AR, DF = "#9aa5b1", "#c2410c"

def ci(p, n): return 1.96 * math.sqrt(p * (1 - p) / n) * 100

fig, ax = plt.subplots(1, 2, figsize=(9, 3.2), gridspec_kw={"width_ratios": [1.35, 1]})
names = ["MMLU\n(n=200)", "GSM8K\n(n=50)"]
ar = [42.0, 21.0]; df = [45.5, 4.0]; ns = [200, 50]
x = np.arange(2)
ax[0].bar(x - .19, ar, .37, color=AR, label="AR abliterated")
ax[0].bar(x + .19, df, .37, color=DF, label="Diffusion")
ax[0].errorbar(x - .19, ar, yerr=[ci(a/100, n) for a, n in zip(ar, ns)], fmt="none", ecolor="#475569", capsize=3)
ax[0].errorbar(x + .19, df, yerr=[ci(d/100, n) for d, n in zip(df, ns)], fmt="none", ecolor="#7c2d12", capsize=3)
ax[0].axhline(25, ls=":", color="#94a3b8", lw=1)
ax[0].text(1.42, 26, "chance", fontsize=7, color="#64748b")
ax[0].set_xticks(x); ax[0].set_xticklabels(names)
ax[0].set_ylabel("accuracy (%)"); ax[0].set_ylim(0, 65)
ax[0].legend(frameon=False, loc="upper right")
ax[0].set_title("Knowledge survives, multi-step reasoning does not", loc="left", fontsize=10)
for i, (a, d) in enumerate(zip(ar, df)):
    ax[0].text(i - .19, a + 2.5, f"{a:.0f}", ha="center", fontsize=7, color="#475569")
    ax[0].text(i + .19, d + 2.5, f"{d:.0f}", ha="center", fontsize=7, color=DF)

sp = [27.3, 61.7, 124.2, 218.9]
lab = ["AR", "64", "32", "16"]
cols = [AR, DF, DF, DF]
ax[1].bar(range(4), sp, color=cols, width=.65)
ax[1].set_xticks(range(4)); ax[1].set_xticklabels(lab)
ax[1].set_xlabel("diffusion steps (AR = baseline)")
ax[1].set_ylabel("tokens/s")
ax[1].set_title("Parallel unmasking wins on throughput", loc="left", fontsize=10)
for i, v in enumerate(sp):
    ax[1].text(i, v + 6, f"{v:.0f}", ha="center", fontsize=7, color="#475569")
ax[1].set_ylim(0, 250)
fig.tight_layout(); fig.savefig(f"{OUT}/fig1_results.png", bbox_inches="tight"); plt.close(fig)

steps = [16, 32, 64]
tps = [218.9, 124.2, 61.7]
fig, ax = plt.subplots(figsize=(6.4, 2.9))
ax.plot(steps, tps, "o-", color=DF, lw=1.6, ms=6, label="diffusion")
ax.axhline(27.3, ls="--", color=AR, lw=1.2, label="AR, KV cache")
ax.set_xlabel("denoising steps for 128 tokens")
ax.set_ylabel("tokens/s")
ax.set_xticks(steps)
ax.legend(frameon=False)
ax.set_title("Throughput is set by the step count", loc="left", fontsize=10)
ax.annotate("8.0x", xy=(16, 218.9), xytext=(24, 205), fontsize=8, color=DF,
            arrowprops=dict(arrowstyle="-", color="#cbd5e1", lw=.8))
ax.annotate("2.3x", xy=(64, 61.7), xytext=(46, 95), fontsize=8, color=DF,
            arrowprops=dict(arrowstyle="-", color="#cbd5e1", lw=.8))
fig.tight_layout(); fig.savefig(f"{OUT}/fig2_speed.png", bbox_inches="tight"); plt.close(fig)

st = [100, 500, 1000, 2000, 3000, 4000]
loss = [0.4137, 0.3258, 0.3482, 0.3147, 0.3018, 0.2953]
fig, ax = plt.subplots(figsize=(6.4, 2.7))
ax.plot(st, loss, "o-", color=DF, lw=1.6, ms=5)
ax.set_xlabel("training step")
ax.set_ylabel("loss (EMA)")
ax.set_ylim(0, 0.5)
ax.set_title("65M tokens, 55 minutes on one H200", loc="left", fontsize=10)
fig.tight_layout(); fig.savefig(f"{OUT}/fig3_loss.png", bbox_inches="tight"); plt.close(fig)
print("ok", os.listdir(OUT))