import matplotlib; matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np

SURFACE, INK, MUTED, GRID = "#faf8f4", "#2c2c2c", "#6b6560", "#d6d0c6"
BLUE, AQUA, YELL, GREEN = "#2a78d6", "#1baf7a", "#eda100", "#008300"

p = np.array([1,2,3,4,5,6,7,8,9])
tpd = np.array([911.3,305.9,155.7,126.1,107.2,136.8,139.6,154.3,145.7])   # ps/DOF measured
fp64 = np.array([7.4,9.4,11.8,15.9,22.4,20.0,21.7,22.3,26.6])
dram = np.array([0.6,1.8,3.4,4.1,4.8,3.8,3.6,3.3,3.4])
shmem= np.array([10.2,24.3,38.0,72.7,84.8,89.5,92.7,93.8,92.7])
occ  = np.array([3.1,3.1,6.2,12.5,21.4,29.8,43.6,57.4,79.9])

# model: bytes/DOF=16, flops/DOF=12(p+1), ridge 4.42 -> flat then rising
FP64_PEAK, BW = 7.737e12, 1751.7e9
fd_model = 12.0*(p+1); ai_model = fd_model/16.0
rate = np.minimum(FP64_PEAK, ai_model*BW)
tpd_model = fd_model/rate*1e12

fig, (axL, axR) = plt.subplots(1, 2, figsize=(11.2, 4.5), dpi=200)
fig.patch.set_facecolor(SURFACE)

# ---------- LEFT: measured vs model ----------
axL.set_facecolor(SURFACE)
axL.plot(p, tpd, color=BLUE, lw=2.2, marker="o", ms=6, mec=SURFACE, mew=1.5, zorder=6)
axL.plot(p, tpd_model, color=MUTED, lw=2, ls=(0,(4,3)), zorder=4)
axL.plot([5],[107.2], "o", ms=13, mfc="none", mec=BLUE, mew=2.2, zorder=7)

axL.text(7.3, 58, "measured", color=BLUE, fontsize=9, fontweight="bold", ha="left")
axL.annotate("roofline model — flat, then rising.\nIt is 12x optimistic: the DRAM bound\nis never approached.",
             xy=(6.6, tpd_model[5]), xytext=(1.6, 26),
             color=MUTED, fontsize=8.3, ha="left", va="center", linespacing=1.4,
             arrowprops=dict(arrowstyle="->", color=MUTED, lw=1.0,
                             connectionstyle="arc3,rad=0.25"))
axL.annotate("sweet spot p = 5\n107 ps/DOF", xy=(5, 107.2), xytext=(5.5, 210),
             color=BLUE, fontsize=8.6, fontweight="bold", ha="left", linespacing=1.35,
             arrowprops=dict(arrowstyle="->", color=BLUE, lw=1.1,
                             connectionstyle="arc3,rad=0.25"))
axL.set_yscale("log")
axL.set_xlabel("polynomial degree  p", color=MUTED, fontsize=9.5)
axL.set_ylabel("time per DOF  (ps, log scale)", color=MUTED, fontsize=9.5)
axL.set_title("Right answer, wrong curve\nThe model predicted p*=5. So did the hardware — by another route.",
              color=INK, fontsize=10.5, loc="left", pad=10, linespacing=1.4)
axL.set_xlim(0.6, 9.5); axL.set_ylim(6, 1600)

# ---------- RIGHT: bottleneck attribution ----------
axR.set_facecolor(SURFACE)
for y, c, name in ((shmem, AQUA, "shared memory"), (occ, GREEN, "occupancy"),
                   (fp64, YELL, "fp64 pipe"), (dram, BLUE, "DRAM")):
    axR.plot(p, y, color=c, lw=2.2, marker="o", ms=5, mec=SURFACE, mew=1.2, zorder=5)
    axR.text(9.22, y[-1], name, color=c, fontsize=8.5, va="center", ha="left", fontweight="bold")
axR.axvline(5, color=INK, lw=1, ls=(0,(2,3)), alpha=0.5, zorder=2)
axR.text(5.15, 63, "p = 5", color=INK, fontsize=8.4, fontweight="bold", va="center")
axR.annotate("shared memory saturates —\nthis is what sets the sweet spot",
             xy=(5, 84.8), xytext=(1.15, 66), color=MUTED, fontsize=8.3, ha="left",
             linespacing=1.4, arrowprops=dict(arrowstyle="->", color=MUTED, lw=1.0,
                                              connectionstyle="arc3,rad=-0.2"))
axR.annotate("DRAM never exceeds 5% of peak —\nthe roofline ridge is never the constraint",
             xy=(7.0, 3.6), xytext=(4.55, 11.5), color=MUTED, fontsize=8.3, ha="left",
             linespacing=1.4, arrowprops=dict(arrowstyle="->", color=MUTED, lw=1.0,
                                              connectionstyle="arc3,rad=-0.25"))
axR.set_xlabel("polynomial degree  p", color=MUTED, fontsize=9.5)
axR.set_ylabel("% of measured peak", color=MUTED, fontsize=9.5)
axR.set_title("Why: the binding resource is on-chip, not off-chip\nA100-SXM4-80GB, fp64, Nsight Compute counters",
              color=INK, fontsize=10.5, loc="left", pad=10, linespacing=1.4)
axR.set_xlim(0.6, 9.5); axR.set_ylim(0, 100)

for ax in (axL, axR):
    ax.grid(True, color=GRID, lw=0.7, alpha=0.75); ax.set_axisbelow(True)
    for s in ("top","right"): ax.spines[s].set_visible(False)
    for s in ("left","bottom"): ax.spines[s].set_color(GRID)
    ax.tick_params(colors=MUTED, labelsize=8.5, length=3)
    ax.set_xticks([1,3,5,7,9])

fig.text(0.008, -0.05,
   "Measured 2026-07-09 on a rented A100-SXM4-80GB (CUDA 12.4). Matrix-free tensor-product Laplacian apply, fp64, sum factorization, one element per\n"
   "block. Machine peaks measured on the same card: fp64 7.74 TFLOP/s, triad 1752 GB/s, ridge 4.42 flop/byte. Measured flops/DOF = 36(p+1)+3 and\n"
   "bytes/DOF ~ 13 -- so arithmetic intensity exceeds the ridge at EVERY degree, including p=1. Naive kernel: 21 syncthreads/element, no element batching.",
   fontsize=7.2, color=MUTED, ha="left", va="top", linespacing=1.5)

fig.tight_layout(rect=(0,0.02,0.995,1))
out = "/Users/srinathvadlamani/Projects/svac/svac-site/static/images/multigrid/ai-sweep-measured.png"
fig.savefig(out, facecolor=SURFACE, bbox_inches="tight", pad_inches=0.28)
print("wrote", out)