#!/usr/bin/env python3
# Replot Fig. 4: demultiplexing-stage complexity ratio vs U/T,
# now including the signed user-wise attention variant.
#
# O_transformer = L * T * d * dff        (FFN-dominated, L = 12 layers)
# O_softmax     = U * d^2 + U^2 * d      (key/value projections + scores)
# O_signed      = U^2 * d + U^2 * 2h + h^2   (Gram + score network, h = 64)
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

T, L, d, h = 32, 12, 128, 64
U = np.arange(16, 65)
x = U / T

plt.rcParams.update({"font.size": 13, "axes.labelsize": 13,
                     "xtick.labelsize": 12, "ytick.labelsize": 12,
                     "axes.linewidth": 1.1, "grid.linewidth": 0.8,
                     "xtick.major.width": 1.1, "ytick.major.width": 1.1,
                     "xtick.minor.width": 0.8, "ytick.minor.width": 0.8,
                     "xtick.major.size": 4.5, "ytick.major.size": 4.5})
fig = plt.figure(figsize=(5.2, 3.9))
ax = fig.add_axes([0.155, 0.145, 0.82, 0.82])
colors = {1.0: "tab:blue", 0.5: "tab:orange", 0.25: "tab:green"}
for r in [1.0, 0.5, 0.25]:                    # r = d / dff
    dff = d / r
    o_tf = L * T * d * dff
    o_soft = U * d ** 2 + U ** 2 * d
    o_sgn = U ** 2 * d + U ** 2 * 2 * h + h ** 2
    ax.plot(x, o_soft / o_tf, color=colors[r], ls="-", lw=1.9,
            label=f"Softmax, $d/d_{{\\mathrm{{ff}}}}$={r:g}")
    ax.plot(x, o_sgn / o_tf, color=colors[r], ls="--", lw=1.9,
            label=f"Signed, $d/d_{{\\mathrm{{ff}}}}$={r:g}")
ax.set_yscale("log")
ax.set_xlabel(r"User-to-token ratio $U/T$")
ax.set_ylabel(r"$\mathcal{O}_{\mathrm{Attention}}/\mathcal{O}_{\mathrm{Transformer}}$")
ax.grid(True, which="both", alpha=0.35)
ax.legend(fontsize=9, ncol=2, loc="lower right")
fig.savefig("fig/complexity_ratio_vs_UT_dff.pdf")
print("saved fig/complexity_ratio_vs_UT_dff.pdf")