"""Generate roofline figures for the paper's Roofline Analysis section.

Produces two 2-panel PNGs into docs/report/1H-codesign-paper/figures/:
  - roofline_short_context.png  (S_kv = 8K)   — batch drops per-token cost
  - roofline_long_context.png   (S_kv = 1M)   — batch effect vanishes

Each figure shows both the step-latency view (raw step time vs B, with
weight/compute/KV components) and the cost-per-token view (÷B), which
together make the batch-vs-context story visible at a glance.

Run: python tests/analytical_visualization/_gen_roofline_paper_figs.py
"""
from __future__ import annotations

import sys
from pathlib import Path

# Ensure the repo root is on sys.path so `tests.analytical_visualization`
# imports resolve when invoked as a script.
_REPO_ROOT = Path(__file__).resolve().parents[2]
if str(_REPO_ROOT) not in sys.path:
    sys.path.insert(0, str(_REPO_ROOT))

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

from tests.analytical_visualization.chip_roofline import (
    critical_batch,
    per_token_latency_curve,
    step_latency_curve,
)
from tests.analytical_visualization.model_config import MachineParams
from tests.analytical_visualization.model_presets import PRESETS

FIGURES_DIR = Path(__file__).resolve().parents[2] / (
    "docs/report/1H-codesign-paper/figures"
)
FIGURES_DIR.mkdir(parents=True, exist_ok=True)

MACHINE = MachineParams()
MODEL = PRESETS["Llama 3 70B"].model
B_RANGE = [1, 2, 4, 8, 16, 32, 64, 128, 256]


def _plot_pair(s_kv: int, label: str, out_path: Path) -> None:
    step = step_latency_curve(MACHINE, MODEL, B_RANGE, s_kv=s_kv)
    tok = per_token_latency_curve(MACHINE, MODEL, B_RANGE, s_kv=s_kv)
    xs = [p.batch for p in step]
    b_star = critical_batch(MACHINE, MODEL)

    fig, (axA, axB) = plt.subplots(1, 2, figsize=(11, 4.2))

    # ── Left: step latency (undivided) ────────────────────────────────
    axA.plot(xs, [p.weight_s * 1e3 for p in step], "^-",
              color="#ffbe0b", label="Weight fetch (flat)")
    axA.plot(xs, [p.compute_s * 1e3 for p in step], "o-",
              color="#3a86ff", label="Compute (linear)")
    axA.plot(xs, [p.kv_s * 1e3 for p in step], "s-",
              color="#d90429", label="KV fetch (linear)")
    axA.plot(xs, [p.total_s * 1e3 for p in step], "-",
              color="#212529", linewidth=2.5, label="Total")
    axA.set_xscale("log", base=2)
    axA.set_yscale("log")
    axA.set_xlabel("Batch size B")
    axA.set_ylabel("Step time (ms)")
    axA.set_title(f"Step latency ({label} context, $S_{{kv}}={s_kv:,}$)")
    axA.grid(True, which="both", alpha=0.3)
    axA.legend(fontsize=8, loc="upper left")

    # ── Right: cost per token (÷ B) ───────────────────────────────────
    axB.plot(xs, [p.weight_s * 1e3 for p in tok], "^-",
              color="#ffbe0b", label="Weight fetch ($\\propto 1/B$)")
    axB.plot(xs, [p.compute_s * 1e3 for p in tok], "o--",
              color="#3a86ff", label="Compute (flat)")
    axB.plot(xs, [p.kv_s * 1e3 for p in tok], "s--",
              color="#d90429", label="KV fetch (flat)")
    axB.plot(xs, [p.total_s * 1e3 for p in tok], "-",
              color="#212529", linewidth=2.5, label="Total")
    axB.axvline(b_star, linestyle=":", color="#2e7d32",
                 label=f"$B^*={b_star:.0f}$")
    axB.set_xscale("log", base=2)
    axB.set_yscale("log")
    axB.set_xlabel("Batch size B")
    axB.set_ylabel("Per-token time (ms)")
    axB.set_title(f"Cost per token ({label} context, $S_{{kv}}={s_kv:,}$)")
    axB.grid(True, which="both", alpha=0.3)
    axB.legend(fontsize=8, loc="upper right")

    plt.tight_layout()
    fig.savefig(out_path, dpi=150, bbox_inches="tight")
    plt.close(fig)
    print(f"  wrote {out_path.name}")


def main() -> None:
    print("Generating roofline figures for the paper ...")
    _plot_pair(8_192, "short", FIGURES_DIR / "roofline_short_context.png")
    _plot_pair(1_048_576, "long", FIGURES_DIR / "roofline_long_context.png")
    print(f"Done. Output: {FIGURES_DIR}")


if __name__ == "__main__":
    main()