"""
Arithmetic-intensity sweep for matrix-free high-order operator evaluation.

ANALYTIC ROOFLINE MODEL — not a hardware measurement.

Model
-----
3D tensor-product (hex) elements, polynomial degree p, n = p+1 nodes/direction,
collocated quadrature (q = n). Laplacian action, fp64.

Matrix-free with sum factorization:
    flops/element  ~= C_sf * d * n^(d+1)      (C_sf = model constant, see sensitivity)
    -> flops/DOF    = C_sf * d * n            (since DOFs/elem = n^d)

    bytes/element   = 8 * n^d * (2 + g)
        2 = read input vector + write output vector
        g = doubles/DOF of *stored* geometry (symmetric 3x3 Jacobian => g=6);
            g=0 means affine or recomputed-on-the-fly geometry.
    -> bytes/DOF    = 8 * (2 + g)             (independent of p)

    => AI(p) = C_sf*d*n / (8*(2+g))           -- LINEAR in p.

Assembled sparse matvec (the thing matrix-free replaces):
    dense element matrix, flops/DOF = 2*n^d, bytes/DOF = 8*n^d
    => AI = 0.25 flop/byte, CONSTANT in p.

Roofline: achievable rate = min(peak_flops, AI * bandwidth)
    time/DOF = (flops/DOF) / rate

Key consequence (derived, then confirmed numerically below):
    - Bandwidth-bound regime (AI < ridge): time/DOF = bytes/DOF / BW = CONSTANT in p.
      Raising p costs NOTHING in time per DOF while buying accuracy per DOF.
    - Compute-bound regime (AI > ridge): time/DOF = flops/DOF / peak, GROWS linearly in p.
    => The sweet spot is the largest p whose AI still sits at/below the machine ridge.

Limitation (stated, not hidden): this model has NO term for register pressure,
shared-memory working set, or occupancy collapse at high p. Those cap real kernels
well below what the roofline alone predicts on high-ridge machines.
"""

D = 3  # spatial dimensions

# Nominal vendor specs, fp64 *vector* (non-tensor-core) peak, and HBM bandwidth.
MACHINES = {
    #                 peak fp64 TFLOP/s, BW GB/s
    "A100-80 SXM":   (9.7,  2039.0),
    "H100 SXM5":     (34.0, 3350.0),
    "MI300X":        (81.7, 5300.0),
}


def ai_matrix_free(p, c_sf, g):
    n = p + 1
    flops_dof = c_sf * D * n
    bytes_dof = 8.0 * (2 + g)
    return flops_dof / bytes_dof, flops_dof, bytes_dof


def ai_assembled(p):
    n = p + 1
    flops_dof = 2.0 * n**D
    bytes_dof = 8.0 * n**D
    return flops_dof / bytes_dof, flops_dof, bytes_dof


def time_per_dof_ns(flops_dof, ai, peak_tflops, bw_gbs):
    peak = peak_tflops * 1e12
    bw = bw_gbs * 1e9
    rate = min(peak, ai * bw)          # roofline
    return flops_dof / rate * 1e9      # nanoseconds per DOF


def ridge(peak_tflops, bw_gbs):
    return (peak_tflops * 1e12) / (bw_gbs * 1e9)


def p_star(c_sf, g, peak_tflops, bw_gbs):
    """Degree at which AI crosses the ridge point (continuous)."""
    r = ridge(peak_tflops, bw_gbs)
    # AI = c_sf*D*(p+1) / (8*(2+g)) = r
    return r * 8.0 * (2 + g) / (c_sf * D) - 1.0


print(__doc__)
print("=" * 108)
print("RIDGE POINTS (fp64 vector peak / HBM bandwidth)")
print("=" * 108)
for m, (pk, bw) in MACHINES.items():
    print(f"  {m:<14}  peak {pk:5.1f} TFLOP/s   BW {bw:6.0f} GB/s   ridge = {ridge(pk,bw):5.2f} flop/byte")

C_SF = 4.0  # baseline model constant
G = 0       # affine / recomputed geometry

print()
print("=" * 108)
print(f"SWEEP: matrix-free, C_sf={C_SF}, geometry g={G} (recomputed).  time/DOF in ns, roofline.")
print("=" * 108)
hdr = f"{'p':>3} {'n':>3} {'flop/DOF':>9} {'byte/DOF':>9} {'AI':>7} |"
for m in MACHINES:
    hdr += f" {m+' t/DOF':>18}"
print(hdr)
print("-" * 108)
for p in range(1, 17):
    ai, fd, bd = ai_matrix_free(p, C_SF, G)
    row = f"{p:>3} {p+1:>3} {fd:>9.0f} {bd:>9.0f} {ai:>7.2f} |"
    for m, (pk, bw) in MACHINES.items():
        t = time_per_dof_ns(fd, ai, pk, bw)
        bound = "C" if ai > ridge(pk, bw) else "B"   # compute- or bandwidth-bound
        row += f" {t:>15.4f} {bound:>2}"
    print(row)

print()
print("  Legend: B = bandwidth-bound (time/DOF flat in p), C = compute-bound (time/DOF grows with p)")
print()
print("=" * 108)
print("SWEET SPOT p* (AI crosses ridge) — sensitivity to model constant C_sf and geometry storage g")
print("=" * 108)
print(f"{'machine':<14} {'ridge':>7} |" + "".join(f" {'C_sf='+str(c):>10}" for c in (2, 4, 6, 8)))
print("-" * 108)
for gg, label in ((0, "g=0  recomputed geometry"), (6, "g=6  stored Jacobian")):
    print(f"  [{label}]")
    for m, (pk, bw) in MACHINES.items():
        row = f"{m:<14} {ridge(pk,bw):>7.2f} |"
        for c in (2, 4, 6, 8):
            row += f" {p_star(c, gg, pk, bw):>10.1f}"
        print(row)
    print()

print("=" * 108)
print("MATRIX-FREE vs ASSEMBLED (the reason any of this works)")
print("=" * 108)
print(f"{'p':>3} {'AI matrix-free':>16} {'AI assembled':>14} {'ratio':>8} {'asm bytes/DOF':>15}")
print("-" * 108)
for p in (1, 2, 4, 8, 16):
    ai_mf, _, _ = ai_matrix_free(p, C_SF, G)
    ai_as, _, bd_as = ai_assembled(p)
    print(f"{p:>3} {ai_mf:>16.2f} {ai_as:>14.2f} {ai_mf/ai_as:>8.1f}x {bd_as:>15.0f}")