"""C6: roofline plot for Qwen3-Coder-30B-A3B on H20.

Reproduces the analytical roofline used in scripts/compute_roofline.py and
plots it as a single PDF: AI vs achievable throughput, with annotated
operating points for prefill at reuse {0, 70, 90, 95}% and decode.

The constants must stay in lockstep with compute_roofline.py. If you change
one, change the other.
"""
import argparse
from pathlib import Path

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

# ---- model constants (mirror scripts/compute_roofline.py) ----
L, D, H_KV, D_HEAD, D_FFN = 48, 2048, 4, 128, 6144
K_EXPERTS = 8
BYTES = 2  # bf16

# ---- H20 ----
PEAK_FLOPS = 148e12
HBM_BW = 4.0e12
RIDGE = PEAK_FLOPS / HBM_BW  # ~37


def attn_prefill_flops(seq_len, new_tokens):
    d_kv = H_KV * D_HEAD
    qkv = new_tokens * (D * D * 2 + D * d_kv * 2 * 2)
    attn = new_tokens * seq_len * D * 2 * 2
    out = new_tokens * D * D * 2
    return (qkv + attn + out) * L


def attn_prefill_bytes(seq_len, new_tokens, cached_tokens):
    d_kv = H_KV * D_HEAD
    weight = D * (D + 2 * d_kv + D) * BYTES * L
    cached_kv = cached_tokens * 2 * d_kv * BYTES * L
    act = new_tokens * D * BYTES * 2 * L
    new_kv = new_tokens * 2 * d_kv * BYTES * L
    return weight + cached_kv + act + new_kv


def ffn_flops(n):
    return 3 * n * D * D_FFN * 2 * K_EXPERTS * L


def ffn_bytes(n):
    weight = K_EXPERTS * 3 * D * D_FFN * BYTES * L
    act = n * D * BYTES * 2 * L
    return weight + act


def point(seq_len, reuse):
    cached = int(seq_len * reuse)
    new = max(1, seq_len - cached)
    f = attn_prefill_flops(seq_len, new) + ffn_flops(new)
    b = attn_prefill_bytes(seq_len, new, cached) + ffn_bytes(new)
    return f, b, new


def decode_point(seq_len):
    f = attn_prefill_flops(seq_len, 1) + ffn_flops(1)
    b = attn_prefill_bytes(seq_len, 1, seq_len) + ffn_bytes(1)
    return f, b


def plot(out_path, seq_len=64000):
    fig, ax = plt.subplots(figsize=(6.5, 4.2))

    ai_grid = np.logspace(-1, 5, 400)
    achievable = np.minimum(ai_grid * HBM_BW, PEAK_FLOPS) / 1e12
    ax.plot(ai_grid, achievable, color="#222", lw=1.5, label="H20 roofline")

    ax.axvline(RIDGE, color="#888", ls=":", lw=1)
    ax.text(RIDGE, 420, f"ridge = {RIDGE:.0f}", color="#666",
            fontsize=8, ha="center", va="top",
            bbox=dict(boxstyle="round,pad=0.2", fc="white", ec="none", alpha=0.85))

    ax.axhline(PEAK_FLOPS / 1e12, color="#aaa", ls="--", lw=0.6)
    ax.text(2, PEAK_FLOPS / 1e12 * 1.08, "compute ceiling (148 TFLOPS bf16)",
            fontsize=8, color="#666", ha="left")

    # operating points: use a legend (not annotations with leader lines, since
    # all 4 prefill points sit on the compute ceiling and would overlap).
    reuses = [0.0, 0.7, 0.9, 0.95]
    colors = ["#d62728", "#ff7f0e", "#2ca02c", "#1f77b4"]
    for reuse, color in zip(reuses, colors):
        f, b, new = point(seq_len, reuse)
        ai = f / b
        thpt = min(ai * HBM_BW, PEAK_FLOPS) / 1e12
        ax.scatter([ai], [thpt], color=color, s=80, zorder=5,
                   edgecolor="white", linewidth=1.2,
                   label=f"prefill reuse={int(reuse*100):>2}%  "
                         f"(new={new:>6,} tok, AI={ai:>6,.0f})")

    f, b = decode_point(seq_len)
    ai_dec = f / b
    thpt_dec = min(ai_dec * HBM_BW, PEAK_FLOPS) / 1e12
    ax.scatter([ai_dec], [thpt_dec], color="#8c564b", s=80, marker="D",
               zorder=5, edgecolor="white", linewidth=1.2,
               label=f"decode  (per-token, seqlen={seq_len:,}, AI={ai_dec:.1f})")

    ax.legend(loc="lower right", fontsize=8.5, framealpha=0.95,
              prop={"family": "monospace", "size": 8})

    ax.set_xscale("log")
    ax.set_yscale("log")
    ax.set_xlim(0.5, 1e5)
    ax.set_ylim(0.5, 500)
    ax.set_xlabel("Arithmetic intensity (FLOP/byte)")
    ax.set_ylabel("Achievable throughput (TFLOPS)")
    ax.set_title(
        f"Prefill stays compute-bound even at 95% reuse  "
        f"(Qwen3-Coder-30B-A3B, H20, seqlen={seq_len:,})",
        fontsize=10,
    )
    ax.grid(True, which="both", alpha=0.25)
    fig.tight_layout()
    fig.savefig(out_path, bbox_inches="tight")
    plt.close(fig)
    print(f"[C6] wrote {out_path}")
    for reuse in reuses:
        f, b, new = point(seq_len, reuse)
        print(f"     reuse={int(reuse*100):>3}%  new={new:>6,}  AI={f/b:>8.1f}  "
              f"bound={'COMPUTE' if f/b > RIDGE else 'MEMORY'}")
    f, b = decode_point(seq_len)
    print(f"     decode             AI={f/b:>8.1f}  bound={'COMPUTE' if f/b > RIDGE else 'MEMORY'}")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--seq-len", type=int, default=64000)
    ap.add_argument("--outdir", default="analysis/pd_sep_paper_section/figures")
    args = ap.parse_args()
    out = Path(args.outdir)
    out.mkdir(parents=True, exist_ok=True)
    plot(out / "fig_c6_roofline.pdf", seq_len=args.seq_len)


if __name__ == "__main__":
    main()