BG   = "#fafaf8"
DARK = "#1a1a1a"

# Color ramp: blue for common tokens, red for rare
TOKEN_COLORS = ["#1a5276", "#2471a3", "#5dade2", "#e67e22", "#c0392b", "#7d2a2a"]

steps = np.arange(N_STEPS)

fig = plt.figure(figsize=(16, 11), facecolor=BG)
fig.suptitle(
    "SGD vs. Adam on Rare Tokens -- Frequency Bias and Variance Normalization",
    fontsize=14, fontweight="bold", color=DARK, y=0.99
)

gs = gridspec.GridSpec(2, 3, figure=fig, hspace=0.45, wspace=0.35)

# ── 1. SGD weight trajectories ────────────────────────────────
ax1 = fig.add_subplot(gs[0, :2])
ax1.set_facecolor(BG)
ax1.axhline(1.0, color=DARK, lw=1, ls="--", alpha=0.3, label="True weight = 1.0")

for i, (token, color) in enumerate(zip(TOKENS, TOKEN_COLORS)):
    ax1.plot(steps, sgd_history[:, i], color=color, lw=1.8,
             label=f"{token} (freq={FREQ[i]:.3f})")

ax1.set_title("SGD -- Weight Trajectories\nRare tokens barely move from zero", fontsize=11, color=DARK)
ax1.set_xlabel("Training Step", fontsize=9)
ax1.set_ylabel("Learned Weight", fontsize=9)
ax1.legend(fontsize=8, loc="right")
ax1.set_ylim(-0.3, 1.6)
ax1.spines[["top", "right"]].set_visible(False)

# Annotate failure zone
ax1.annotate(
    "Rare tokens stuck\nnear zero",
    xy=(N_STEPS * 0.95, sgd_history[-1, 5]),
    xytext=(N_STEPS * 0.65, -0.15),
    fontsize=8.5, color="#c0392b",
    arrowprops=dict(arrowstyle="->", color="#c0392b", lw=1.2),
    bbox=dict(boxstyle="round,pad=0.3", facecolor="#fff0f0", edgecolor="#c0392b", alpha=0.85)
)

# ── 2. Final weight error bar chart ───────────────────────────
ax2 = fig.add_subplot(gs[0, 2])
ax2.set_facecolor(BG)

x      = np.arange(6)
w_sgd  = sgd_final
w_adam = adam_final
width  = 0.35

bars_sgd  = ax2.bar(x - width/2, np.abs(w_sgd  - TRUE_W), width, color="#c0392b", alpha=0.85, label="SGD error")
bars_adam = ax2.bar(x + width/2, np.abs(w_adam - TRUE_W), width, color="#2980b9", alpha=0.85, label="Adam error")

ax2.set_xticks(x)
ax2.set_xticklabels([t[:8] for t in TOKENS], rotation=30, ha="right", fontsize=8)
ax2.set_ylabel("|learned w − true w|", fontsize=9)
ax2.set_title("Final Weight Error\n(lower = better)", fontsize=11, color=DARK)
ax2.legend(fontsize=8)
ax2.spines[["top", "right"]].set_visible(False)

# ── 3. Adam weight trajectories ───────────────────────────────
ax3 = fig.add_subplot(gs[1, :2])
ax3.set_facecolor(BG)
ax3.axhline(1.0, color=DARK, lw=1, ls="--", alpha=0.3, label="True weight = 1.0")

for i, (token, color) in enumerate(zip(TOKENS, TOKEN_COLORS)):
    ax3.plot(steps, adam_history[:, i], color=color, lw=1.8,
             label=f"{token} (freq={FREQ[i]:.3f})")

ax3.set_title("Adam -- Weight Trajectories\nRare tokens converge via variance normalization", fontsize=11, color=DARK)
ax3.set_xlabel("Training Step", fontsize=9)
ax3.set_ylabel("Learned Weight", fontsize=9)
ax3.legend(fontsize=8, loc="right")
ax3.set_ylim(-0.3, 1.6)
ax3.spines[["top", "right"]].set_visible(False)

ax3.annotate(
    "Rare tokens converge\ndespite sparse gradients",
    xy=(N_STEPS * 0.95, adam_history[-1, 5]),
    xytext=(N_STEPS * 0.60, 0.3),
    fontsize=8.5, color="#27ae60",
    arrowprops=dict(arrowstyle="->", color="#27ae60", lw=1.2),
    bbox=dict(boxstyle="round,pad=0.3", facecolor="#f0fff4", edgecolor="#27ae60", alpha=0.85)
)

# ── 4. Effective LR vs frequency ─────────────────────────────
ax4 = fig.add_subplot(gs[1, 2])
ax4.set_facecolor(BG)

ax4.scatter(FREQ, effective_lr, c=TOKEN_COLORS, s=120, zorder=5, edgecolors="white", lw=1.5)
for i, token in enumerate(TOKENS):
    ax4.annotate(token, (FREQ[i], effective_lr[i]),
                 textcoords="offset points", xytext=(6, 4), fontsize=7.5, color=TOKEN_COLORS[i])

ax4.axhline(LR, color=DARK, lw=1, ls="--", alpha=0.4)
ax4.text(0.5, LR * 1.05, f"Nominal LR = {LR}", fontsize=8, color=DARK, alpha=0.6)

ax4.set_xscale("log")
ax4.set_yscale("log")
ax4.set_xlabel("Token Frequency (log scale)", fontsize=9)
ax4.set_ylabel("Adam Effective LR  lr/√v̂  (log scale)", fontsize=9)
ax4.set_title("Adam's Automatic Equalizer\nRare tokens get amplified LR", fontsize=11, color=DARK)
ax4.spines[["top", "right"]].set_visible(False)

plt.savefig("sgd_vs_adam.png", dpi=150, bbox_inches="tight", facecolor=BG)
plt.show()