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

# --- site tokens (warm parchment) ---
SURFACE = "#faf8f4"
INK     = "#2c2c2c"
MUTED   = "#6b6560"
GRID    = "#d6d0c6"
# --- validated categorical slots 1,2,3 ---
BLUE, AQUA, YELL = "#2a78d6", "#1baf7a", "#eda100"

D, C_SF = 3, 4.0
MACH = [("A100-80", 9.7, 2039.0, BLUE),
        ("H100",   34.0, 3350.0, AQUA),
        ("MI300X", 81.7, 5300.0, YELL)]

p = np.arange(1, 17)
n = p + 1

def ai(g):   return C_SF * D * n / (8.0 * (2 + g))
def fd():    return C_SF * D * n                      # flops/DOF
def tpd(peak_t, bw_g, g):
    a = ai(g)
    rate = np.minimum(peak_t * 1e12, a * bw_g * 1e9)
    return fd() / rate * 1e12                          # picoseconds / DOF

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

# ============ LEFT: AI(p) vs ridge points ============
axL.set_facecolor(SURFACE)
axL.plot(p, ai(0), color=INK, lw=2, zorder=5)
axL.plot(p, ai(6), color=MUTED, lw=2, ls=(0, (4, 3)), zorder=4)

axL.annotate("matrix-free\ngeometry recomputed", xy=(11.4, ai(0)[10]), xytext=(9.1, 11.6),
             color=INK, fontsize=8.5, ha="left", va="bottom", linespacing=1.35)
axL.annotate("stored Jacobian —\n4x the bytes/DOF, 4x the time", xy=(10.6, ai(6)[9]), xytext=(16.0, 0.25),
             color=MUTED, fontsize=8.5, ha="right", va="bottom", linespacing=1.35,
             arrowprops=dict(arrowstyle="->", color=MUTED, lw=1.0,
                             connectionstyle="arc3,rad=0.25"))

for name, pk, bw, col in MACH:
    r = pk * 1e12 / (bw * 1e9)
    axL.axhline(r, color=col, lw=1.6, ls=(0, (1, 2.2)), zorder=3)
    axL.text(16.35, r, f"{name}\nridge {r:.1f}", color=col, fontsize=8.2,
             va="center", ha="left", linespacing=1.3, fontweight="bold")
    ps = r * 8.0 * 2 / (C_SF * D) - 1.0
    if ps <= 16:
        axL.plot([ps], [r], "o", ms=7, color=col, mec=SURFACE, mew=1.8, zorder=6)

axL.set_xlabel("polynomial degree  p", color=MUTED, fontsize=9.5)
axL.set_ylabel("arithmetic intensity  (flop / byte)", color=MUTED, fontsize=9.5)
axL.set_title("Intensity rises linearly with order—\nand each machine's ridge sets where that stops paying",
              color=INK, fontsize=10.5, loc="left", pad=10, linespacing=1.4)
axL.set_xlim(1, 16); axL.set_ylim(0, 17)

# ============ RIGHT: time/DOF -> the sweet spot ============
axR.set_facecolor(SURFACE)
for name, pk, bw, col in MACH:
    t = tpd(pk, bw, 0)
    axR.plot(p, t, color=col, lw=2, zorder=5)
    axR.text(16.3, t[-1], name, color=col, fontsize=8.6, va="center",
             ha="left", fontweight="bold")
    r = pk * 1e12 / (bw * 1e9)
    ps = r * 8.0 * 2 / (C_SF * D) - 1.0
    if ps <= 16:
        axR.plot([ps], [t[0]], "o", ms=7, color=col, mec=SURFACE, mew=1.8, zorder=6)

axR.axvspan(1, 5.3, color=BLUE, alpha=0.055, lw=0, zorder=1)
axR.text(3.1, 19.3, "free accuracy\n(bandwidth-bound:\ntime/DOF is flat)", color=MUTED,
         fontsize=8.3, ha="center", va="top", linespacing=1.4)
axR.annotate("past the ridge every extra\norder costs time",
             xy=(13.0, 17.3), xytext=(7.5, 20.6),
             color=MUTED, fontsize=8.3, ha="left", va="top", linespacing=1.4,
             arrowprops=dict(arrowstyle="->", color=MUTED, lw=1.1,
                             connectionstyle="arc3,rad=-0.25"))
axR.annotate("p* = 5", xy=(5.3, 7.85), xytext=(5.9, 10.2), color=BLUE, fontsize=8.6,
             fontweight="bold", ha="left", va="bottom")

axR.set_xlabel("polynomial degree  p", color=MUTED, fontsize=9.5)
axR.set_ylabel("time per DOF  (ps, roofline)", color=MUTED, fontsize=9.5)
axR.set_title("The sweet spot: raise p until intensity hits the ridge,\nthen stop",
              color=INK, fontsize=10.5, loc="left", pad=10, linespacing=1.4)
axR.set_xlim(1, 16); axR.set_ylim(0, 23)

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, 4, 8, 12, 16])

fig.text(0.008, -0.055,
         "Analytic roofline model, not a hardware measurement.  3D tensor-product elements, fp64, sum factorization "
         "(flops/DOF = C·d·(p+1), C=4);\nbytes/DOF = 8·(2+g) is independent of p.  Nominal fp64 vector peak and HBM "
         "bandwidth.  No register-pressure or occupancy term, so p* is an\nupper bound on high-ridge machines.",
         fontsize=7.3, 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.png"
import os; os.makedirs(os.path.dirname(out), exist_ok=True)
fig.savefig(out, facecolor=SURFACE, bbox_inches="tight", pad_inches=0.28)
print("wrote", out)