"""Per-stage cost formulas for attention forward pass.

Each stage returns:
    name, formula (str with the substituted numeric form),
    compute_time, memory_time, comm_time, bound, visible_time.

Bound is the max of the three; "hidden" comm can appear as 0 when
overlapped with compute.

All times are in seconds; convert to μs at display.
"""
from __future__ import annotations

from dataclasses import dataclass
from typing import Literal

from .model_config import FullConfig


Bound = Literal["compute", "memory", "comm", "trivial"]


@dataclass
class StageCost:
    name: str
    formula: str
    compute_s: float
    memory_s: float
    comm_s: float
    bound: Bound
    visible_s: float
    hop_multiplier: int = 1  # some stages scale with N_cp hops
    # Detailed breakdown for the per-stage table.
    flops: int = 0                  # total FLOPs (numeric)
    mem_bytes: int = 0              # HBM bytes moved (numeric)
    comm_bytes: int = 0             # comm bytes over the wire (numeric)
    flops_formula: str = ""         # human-readable FLOPs formula
    mem_formula: str = ""           # human-readable memory formula
    comm_formula: str = ""          # human-readable comm formula


def _visible(compute: float, memory: float, comm: float) -> tuple[float, Bound]:
    """Pick the dominant time as the visible stage cost."""
    parts = {"compute": compute, "memory": memory, "comm": comm}
    dominant = max(parts, key=parts.get)  # type: ignore
    return parts[dominant], dominant  # type: ignore


def stage_rmsnorm(cfg: FullConfig) -> StageCost:
    T_q = cfg.topo.T_q
    d = cfg.model.hidden
    b = cfg.model.bytes_per_elem
    bytes_ = T_q * d * b * 2  # load x + load weight
    flops = 4 * T_q * d
    mem_s = bytes_ / cfg.machine.bw_hbm
    cmp_s = flops / (cfg.machine.peak_flops * cfg.machine.compute_util)
    vis, bnd = _visible(cmp_s, mem_s, 0)
    return StageCost(
        name="S1 RMSNorm",
        formula=f"bytes = {T_q}*{d}*{b}*2 = {bytes_} B / BW_HBM",
        compute_s=cmp_s, memory_s=mem_s, comm_s=0.0,
        bound=bnd, visible_s=vis,
        flops=flops, mem_bytes=bytes_,
        flops_formula=f"4*T_q*d = 4*{T_q}*{d} = {flops}",
        mem_formula=f"2*T_q*d*b = 2*{T_q}*{d}*{b} = {bytes_} B",
    )


def _gemm_time(flops: int, weight_bytes: int, cfg: FullConfig) -> tuple[float, float]:
    cmp_s = flops / (cfg.machine.peak_flops * cfg.machine.compute_util)
    mem_s = weight_bytes / cfg.machine.bw_hbm
    return cmp_s, mem_s


def stage_wq(cfg: FullConfig) -> StageCost:
    T_q = cfg.topo.T_q
    d = cfg.model.hidden
    b = cfg.model.bytes_per_elem
    hq_per_pe = cfg.h_q_per_pe
    dh = cfg.model.d_head
    flops = 2 * T_q * d * (hq_per_pe * dh)
    weight_B = d * (hq_per_pe * dh) * b
    cmp_s, mem_s = _gemm_time(flops, weight_B, cfg)
    vis, bnd = _visible(cmp_s, mem_s, 0)
    return StageCost(
        name="S2 W_Q GEMM",
        formula=f"FLOPs = 2*{T_q}*{d}*{hq_per_pe*dh} = {flops:.2g}; "
                f"weight = {weight_B/1e6:.1f} MB",
        compute_s=cmp_s, memory_s=mem_s, comm_s=0.0,
        bound=bnd, visible_s=vis,
        flops=int(flops), mem_bytes=int(weight_B),
        flops_formula=f"2*T_q*d*(H_q/TP*d_h) = 2*{T_q}*{d}*{hq_per_pe*dh} = {flops:.3g}",
        mem_formula=f"d*(H_q/TP*d_h)*b = {d}*{hq_per_pe*dh}*{b} = {weight_B/1e6:.2f} MB",
    )


def stage_wkv(cfg: FullConfig) -> StageCost:
    T_q = cfg.topo.T_q
    d = cfg.model.hidden
    b = cfg.model.bytes_per_elem
    hkv_per_pe = max(1, cfg.model.h_kv // cfg.topo.tp)
    dh = cfg.model.d_head
    flops_one = 2 * T_q * d * (hkv_per_pe * dh)
    weight_B_one = d * (hkv_per_pe * dh) * b
    flops = 2 * flops_one
    weight_B = 2 * weight_B_one
    cmp_s, mem_s = _gemm_time(flops, weight_B, cfg)
    vis, bnd = _visible(cmp_s, mem_s, 0)
    return StageCost(
        name="S3 W_K + W_V GEMM",
        formula=f"FLOPs per proj = 2*{T_q}*{d}*{hkv_per_pe*dh} = {flops_one:.2g}; x2",
        compute_s=cmp_s, memory_s=mem_s, comm_s=0.0,
        bound=bnd, visible_s=vis,
        flops=int(flops), mem_bytes=int(weight_B),
        flops_formula=f"2*(2*T_q*d*(H_kv/TP*d_h)) = 2*(2*{T_q}*{d}*{hkv_per_pe*dh}) = {flops:.3g}",
        mem_formula=f"2*(d*(H_kv/TP*d_h)*b) = 2*({d}*{hkv_per_pe*dh}*{b}) = {weight_B/1e6:.2f} MB",
    )


def stage_kv_append(cfg: FullConfig) -> StageCost:
    T_q = cfg.topo.T_q
    b = cfg.model.bytes_per_elem
    hkv_per_pe = max(1, cfg.model.h_kv // cfg.topo.tp)
    dh = cfg.model.d_head
    bytes_ = 2 * T_q * hkv_per_pe * dh * b
    mem_s = bytes_ / cfg.machine.bw_hbm
    return StageCost(
        name="S4 KV cache append",
        formula=f"bytes = 2*{T_q}*{hkv_per_pe}*{dh}*{b} = {bytes_} B",
        compute_s=0, memory_s=mem_s, comm_s=0,
        bound="memory", visible_s=mem_s,
        flops=0, mem_bytes=bytes_,
        flops_formula="0 (no matmul, just cache write)",
        mem_formula=f"2*T_q*(H_kv/TP)*d_h*b = 2*{T_q}*{hkv_per_pe}*{dh}*{b} = {bytes_} B",
    )


def _per_hop_qkT_pv(cfg: FullConfig) -> tuple[float, str]:
    """Q·Kᵀ (and P·V has same FLOPs) per hop, per PE."""
    T_q = cfg.topo.T_q
    S_local = cfg.topo.s_local
    dh = cfg.model.d_head
    hq_per_pe = cfg.h_q_per_pe
    flops = 2 * T_q * S_local * dh * hq_per_pe
    cmp_s = flops / (cfg.machine.peak_flops * cfg.machine.compute_util)
    formula = f"2*{T_q}*{S_local}*{dh}*{hq_per_pe} = {flops:.2g} FLOPs/hop"
    return cmp_s, formula


def _cp_compute_passes(cfg: FullConfig) -> int:
    """How many local S5/S6/S7 passes happen per token step under CP.
    - decode: 1 (each PE computes once against its local K,V, then all-reduce O/m/l)
    - prefill: CP  (K/V or Q rotates through the ring; one local pass per hop)
    """
    if cfg.topo.cp <= 1:
        return 1
    return 1 if cfg.topo.mode == "decode" else cfg.topo.cp


def stage_qkT(cfg: FullConfig) -> StageCost:
    cmp_hop, _ = _per_hop_qkT_pv(cfg)
    passes = _cp_compute_passes(cfg)
    total = cmp_hop * passes
    T_q = cfg.topo.T_q
    S_local = cfg.topo.s_local
    dh = cfg.model.d_head
    hq_per_pe = cfg.h_q_per_pe
    flops_per_hop = 2 * T_q * S_local * dh * hq_per_pe
    total_flops = flops_per_hop * passes
    _hop_word = "pass" if passes == 1 else "hops"
    return StageCost(
        name=f"S5 Q.K^T (x{passes} {_hop_word})",
        formula=f"2*{T_q}*{S_local}*{dh}*{hq_per_pe} = {flops_per_hop:.2g} FLOPs/hop",
        compute_s=total, memory_s=0, comm_s=0,
        bound="compute", visible_s=total,
        hop_multiplier=passes,
        flops=int(total_flops), mem_bytes=0,
        flops_formula=(
            f"{passes}*(2*T_q*S_local*d_h*(H_q/TP)) = "
            f"{passes}*(2*{T_q}*{S_local}*{dh}*{hq_per_pe}) = {total_flops:.3g}"
        ),
        mem_formula="0 (scores accumulated on-chip)",
    )


def stage_softmax(cfg: FullConfig) -> StageCost:
    T_q = cfg.topo.T_q
    S_local = cfg.topo.s_local
    hq_per_pe = cfg.h_q_per_pe
    b = cfg.model.bytes_per_elem
    elems = hq_per_pe * T_q * S_local
    bytes_ = elems * b * 2
    mem_s_per_hop = bytes_ / cfg.machine.bw_hbm
    passes = _cp_compute_passes(cfg)
    total = mem_s_per_hop * passes
    total_bytes = bytes_ * passes
    _hop_word = "pass" if passes == 1 else "hops"
    return StageCost(
        name=f"S6 softmax (x{passes} {_hop_word})",
        formula=f"elts/hop = {hq_per_pe}*{T_q}*{S_local} = {elems:.2g}",
        compute_s=0, memory_s=total, comm_s=0,
        bound="memory", visible_s=total,
        hop_multiplier=passes,
        flops=0, mem_bytes=int(total_bytes),
        flops_formula="~O(elts) (negligible)",
        mem_formula=(
            f"{passes}*2*b*(H_q/TP)*T_q*S_local = "
            f"{passes}*2*{b}*{hq_per_pe}*{T_q}*{S_local} = "
            f"{total_bytes/1e6:.2f} MB"
        ),
    )


def stage_pv(cfg: FullConfig) -> StageCost:
    cmp_hop, _ = _per_hop_qkT_pv(cfg)
    passes = _cp_compute_passes(cfg)
    total = cmp_hop * passes
    T_q = cfg.topo.T_q
    S_local = cfg.topo.s_local
    dh = cfg.model.d_head
    hq_per_pe = cfg.h_q_per_pe
    flops_per_hop = 2 * T_q * S_local * dh * hq_per_pe
    total_flops = flops_per_hop * passes
    _hop_word = "pass" if passes == 1 else "hops"
    return StageCost(
        name=f"S7 P.V (x{passes} {_hop_word})",
        formula=f"2*{T_q}*{S_local}*{dh}*{hq_per_pe} = {flops_per_hop:.2g} FLOPs/hop",
        compute_s=total, memory_s=0, comm_s=0,
        bound="compute", visible_s=total,
        hop_multiplier=passes,
        flops=int(total_flops), mem_bytes=0,
        flops_formula=(
            f"{passes}*(2*T_q*S_local*d_h*(H_q/TP)) = "
            f"{passes}*(2*{T_q}*{S_local}*{dh}*{hq_per_pe}) = {total_flops:.3g}"
        ),
        mem_formula="0 (accumulated on-chip)",
    )


def stage_merge(cfg: FullConfig) -> StageCost:
    """S8: online-softmax merge of partial (o, m, l).
    - prefill (K/V ring): merge happens in-place across CP hops; no comm here.
    - decode: merge is the moment we all-reduce partial (o, m, l) across CP
      ranks; comm is folded into this stage (no separate C1).
    """
    T_q = cfg.topo.T_q
    hq_per_pe = cfg.h_q_per_pe
    dh = cfg.model.d_head
    b = cfg.model.bytes_per_elem
    flops = 6 * T_q * hq_per_pe * dh * max(0, cfg.topo.cp - 1)
    cmp_s = flops / (cfg.machine.peak_flops * cfg.machine.compute_util)

    comm_s = 0.0
    comm_bytes = 0
    comm_formula = ""
    name_suffix = ""

    if cfg.topo.mode == "decode" and cfg.topo.cp > 1:
        cp = cfg.topo.cp
        # (O + m + l) bytes per rank
        M = (T_q * hq_per_pe * dh * b) + (2 * T_q * hq_per_pe * b)
        if cfg.topo.cp_placement == "pe":
            bw, alpha, tier = cfg.machine.bw_intra, cfg.machine.alpha_intra, "intra-cube"
        elif cfg.topo.sips_used > 1:
            bw, alpha, tier = cfg.machine.bw_intersip, cfg.machine.alpha_intersip, "inter-SIP"
        else:
            bw, alpha, tier = cfg.machine.bw_inter, cfg.machine.alpha_inter, "inter-cube"
        ar_bytes = 2 * (cp - 1) / cp * M
        comm_s = ar_bytes / bw + 2 * (cp - 1) * alpha
        comm_bytes = int(ar_bytes)
        comm_formula = (
            f"PURPOSE (folded into S8 in decode): each CP rank computed\n"
            f"attention against its LOCAL slice of KV and holds a partial\n"
            f"(O, m, l). This all-reduce merges those partials across all\n"
            f"{cp} CP ranks (online-softmax combine). No separate C1 row\n"
            f"in decode because this is the only CP comm and it happens\n"
            f"exactly here, at the end of the local attention body.\n"
            f"---\n"
            f"AR of (O + m + l): 2*(CP-1)/CP * M "
            f"= 2*({cp}-1)/{cp} * {M} B "
            f"= {ar_bytes:.0f} B over {tier} ({bw/1e9:.0f} GB/s) "
            f"+ 2*({cp}-1)*alpha"
        )
        name_suffix = f" + O/m/l AR (CP={cp} ranks, {tier})"

    visible_s = max(cmp_s, comm_s)
    if comm_s > cmp_s:
        bound = "comm"
    elif cmp_s > 0:
        bound = "compute"
    else:
        bound = "trivial"

    return StageCost(
        name=f"S8 online-softmax merge (x{max(0, cfg.topo.cp-1)}){name_suffix}",
        formula=f"~6*{T_q}*{hq_per_pe}*{dh}*(C-1) = {flops:.2g} FLOPs",
        compute_s=cmp_s, memory_s=0, comm_s=comm_s,
        bound=bound, visible_s=visible_s,
        flops=int(flops), mem_bytes=0, comm_bytes=comm_bytes,
        flops_formula=(
            f"6*T_q*(H_q/TP)*d_h*(CP-1) = "
            f"6*{T_q}*{hq_per_pe}*{dh}*{max(0, cfg.topo.cp-1)} = {flops}"
        ),
        mem_formula="0 (in-register)",
        comm_formula=comm_formula or "0",
    )


def stage_normalize(cfg: FullConfig) -> StageCost:
    T_q = cfg.topo.T_q
    hq_per_pe = cfg.h_q_per_pe
    dh = cfg.model.d_head
    flops = T_q * hq_per_pe * dh
    cmp_s = flops / (cfg.machine.peak_flops * cfg.machine.compute_util)
    return StageCost(
        name="S9 normalize O/l",
        formula=f"{T_q}*{hq_per_pe}*{dh} = {flops} divisions",
        compute_s=cmp_s, memory_s=0, comm_s=0,
        bound="trivial", visible_s=cmp_s,
        flops=int(flops), mem_bytes=0,
        flops_formula=f"T_q*(H_q/TP)*d_h = {T_q}*{hq_per_pe}*{dh} = {flops}",
        mem_formula="0",
    )


def stage_wo(cfg: FullConfig) -> StageCost:
    T_q = cfg.topo.T_q
    d = cfg.model.hidden
    b = cfg.model.bytes_per_elem
    hq_per_pe = cfg.h_q_per_pe
    dh = cfg.model.d_head
    flops = 2 * T_q * (hq_per_pe * dh) * d
    weight_B = (hq_per_pe * dh) * d * b
    cmp_s, mem_s = _gemm_time(flops, weight_B, cfg)
    vis, bnd = _visible(cmp_s, mem_s, 0)
    return StageCost(
        name="S10 W_O GEMM",
        formula=f"FLOPs = 2*{T_q}*{hq_per_pe*dh}*{d} = {flops:.2g}; "
                f"weight = {weight_B/1e6:.1f} MB",
        compute_s=cmp_s, memory_s=mem_s, comm_s=0.0,
        bound=bnd, visible_s=vis,
        flops=int(flops), mem_bytes=int(weight_B),
        flops_formula=f"2*T_q*(H_q/TP*d_h)*d = 2*{T_q}*{hq_per_pe*dh}*{d} = {flops:.3g}",
        mem_formula=f"(H_q/TP*d_h)*d*b = {hq_per_pe*dh}*{d}*{b} = {weight_B/1e6:.2f} MB",
    )


def comm_cp_ring(cfg: FullConfig) -> StageCost:
    """CP comm — depends on mode:

    - decode: single O/m/l all-reduce over CP ranks AFTER local S5-S8 finish.
      No per-hop ring; each PE computes attention against its local K,V once
      and then contributes (o, m, l) to the reduce.
    - prefill: per-hop ring during S5/S7 compute (either K/V or Q+O/m/l per
      cp_ring_variant).

    BW tier depends on cp_placement:
    - cp_placement=pe:   intra-cube BW for all hops
    - cp_placement=cube: inter-cube BW; long rings cross SIP boundaries
    """
    if cfg.topo.cp <= 1:
        return StageCost(
            name="C1 CP comm", formula="C=1 -> no comm",
            compute_s=0, memory_s=0, comm_s=0, bound="trivial", visible_s=0,
        )
    S_local = cfg.topo.s_local
    hkv_per_pe = max(1, cfg.model.h_kv // cfg.topo.tp)
    hq_per_pe = cfg.h_q_per_pe
    T_q = cfg.topo.T_q
    dh = cfg.model.d_head
    b = cfg.model.bytes_per_elem

    # ── Decode: single O/m/l all-reduce after local compute ──
    if cfg.topo.mode == "decode":
        cp = cfg.topo.cp
        # per-rank O + m + l bytes
        M = (T_q * hq_per_pe * dh * b) + (2 * T_q * hq_per_pe * b)
        if cfg.topo.cp_placement == "pe":
            bw = cfg.machine.bw_intra
            alpha = cfg.machine.alpha_intra
            tier = "intra-cube"
        elif cfg.topo.sips_used > 1:
            bw = cfg.machine.bw_intersip
            alpha = cfg.machine.alpha_intersip
            tier = "inter-SIP"
        else:
            bw = cfg.machine.bw_inter
            alpha = cfg.machine.alpha_inter
            tier = "inter-cube"
        ar_bytes = 2 * (cp - 1) / cp * M
        ar_time = ar_bytes / bw + 2 * (cp - 1) * alpha
        return StageCost(
            name=f"C1 O/m/l all-reduce ({tier}, x{cp} ranks)",
            formula=(f"decode: gather partial (O,m,l) once at end; "
                     f"M = T_q*(H_q/TP)*(d_h*b + 2*b) = {M} B per rank; "
                     f"AR: 2*(CP-1)/CP*M / BW + 2*(CP-1)*alpha"),
            compute_s=0, memory_s=0, comm_s=ar_time,
            bound="comm", visible_s=ar_time,
            flops=0, mem_bytes=0, comm_bytes=int(ar_bytes),
            flops_formula="0",
            mem_formula="0",
            comm_formula=(
                f"PURPOSE: In decode, each CP rank computed attention against\n"
                f"its OWN slice of the KV cache. Each rank now holds a "
                f"partial\n"
                f"(O, m, l). This single all-reduce merges those partials "
                f"across\n"
                f"all {cp} CP ranks (using online-softmax math) to get the "
                f"final O.\n"
                f"---\n"
                f"AR of (O + m + l): 2*(CP-1)/CP * M "
                f"= 2*({cp}-1)/{cp} * {M} "
                f"= {ar_bytes:.0f} B over {tier} at {bw/1e9:.0f} GB/s "
                f"+ 2*({cp}-1)*alpha"
            ),
        )

    # ── Prefill: per-hop ring (K/V or Q+O/m/l) concurrent with S5-S7 ──
    if cfg.topo.cp_ring_variant == "qoml":
        M_KV = (2 * T_q * hq_per_pe * dh * b) + (2 * T_q * hq_per_pe * b)
        variant_desc = "Q+O/m/l ring"
        formula_bytes = (
            f"2*T_q*(H_q/TP)*d_h*b + 2*T_q*(H_q/TP)*b "
            f"= 2*{T_q}*{hq_per_pe}*{dh}*{b} + 2*{T_q}*{hq_per_pe}*{b} "
            f"= {M_KV} B/hop"
        )
    else:
        M_KV = 2 * S_local * hkv_per_pe * dh * b
        variant_desc = "K/V ring"
        formula_bytes = (
            f"2*S_local*(H_kv/TP)*d_h*b "
            f"= 2*{S_local}*{hkv_per_pe}*{dh}*{b} "
            f"= {M_KV/1e6:.3f} MB/hop"
        )

    if cfg.topo.cp_placement == "pe":
        # All hops intra-cube (fastest).
        intra_hops = cfg.topo.cp - 1
        inter_hops = 0
        intra_time = (intra_hops * M_KV / cfg.machine.bw_intra
                      + intra_hops * cfg.machine.alpha_intra)
        inter_time = 0.0
    else:  # cp_placement == "cube"
        intra_hops = cfg.topo.cp_intra_sip_hops
        inter_hops = cfg.topo.cp_inter_sip_hops
        intra_time = (intra_hops * M_KV / cfg.machine.bw_inter
                      + intra_hops * cfg.machine.alpha_inter)
        inter_time = (inter_hops * M_KV / cfg.machine.bw_intersip
                      + inter_hops * cfg.machine.alpha_intersip)
    comm_time = intra_time + inter_time

    # Overlap check: per-hop compute (S5+S6+S7) hides intra-SIP hops well;
    # inter-SIP hops usually dominate.
    per_hop_cmp, _ = _per_hop_qkT_pv(cfg)
    per_hop_mem_softmax = (cfg.h_q_per_pe * cfg.topo.T_q
                          * S_local * cfg.model.bytes_per_elem * 2
                          / cfg.machine.bw_hbm)
    per_hop_compute = 2 * per_hop_cmp + per_hop_mem_softmax
    per_intra_ring = M_KV / cfg.machine.bw_inter + cfg.machine.alpha_inter
    per_inter_ring = M_KV / cfg.machine.bw_intersip + cfg.machine.alpha_intersip
    visible_intra = intra_hops * max(0.0, per_intra_ring - per_hop_compute)
    visible_inter = inter_hops * max(0.0, per_inter_ring - per_hop_compute)
    visible_total = visible_intra + visible_inter

    tier_desc = f"{intra_hops} intra-SIP + {inter_hops} inter-SIP hops"
    total_ring_bytes = M_KV * (intra_hops + inter_hops)
    _var_purpose = (
        "K, V shards rotate between CP ranks each hop"
        if cfg.topo.cp_ring_variant == "kv"
        else "Q + running (O, m, l) rotate between CP ranks each hop"
    )
    return StageCost(
        name=f"C1 CP {variant_desc} ({tier_desc})",
        formula=f"{formula_bytes}; "
                f"intra: {intra_hops}*M/BW + inter: {inter_hops}*M/BW",
        compute_s=0, memory_s=0, comm_s=comm_time,
        bound="comm", visible_s=visible_total,
        flops=0, mem_bytes=0, comm_bytes=int(total_ring_bytes),
        flops_formula="0",
        mem_formula="0",
        comm_formula=(
            f"PURPOSE: CP shards the sequence axis. Each CP rank holds only\n"
            f"1/{cfg.topo.cp} of the KV cache, so to compute full attention\n"
            f"we must move data between CP ranks. Variant: {_var_purpose}.\n"
            f"Runs concurrently with S5/S6/S7 - each hop, compute of the\n"
            f"just-arrived shard overlaps with comm of the next shard.\n"
            f"---\n"
            f"M*(CP-1) with M = {formula_bytes}; "
            f"total = {total_ring_bytes/1e6:.3f} MB over the ring"
        ),
    )


def comm_tp_allreduce(cfg: FullConfig) -> StageCost:
    """TP AllReduce on W_O output."""
    if cfg.topo.tp <= 1:
        return StageCost(
            name="C2 TP AllReduce W_O", formula="TP=1 -> no AR",
            compute_s=0, memory_s=0, comm_s=0, bound="trivial", visible_s=0,
        )
    T_q = cfg.topo.T_q
    d = cfg.model.hidden
    b = cfg.model.bytes_per_elem
    bytes_ = T_q * d * b
    tp = cfg.topo.tp
    tier = cfg.topo.tp_link_tier()  # "intra" | "inter" | "intersip"
    if tier == "intra":
        bw, alpha, scope = cfg.machine.bw_intra, cfg.machine.alpha_intra, "intra-cube"
    elif tier == "inter":
        bw, alpha, scope = cfg.machine.bw_inter, cfg.machine.alpha_inter, "cross-cube"
    else:
        bw, alpha, scope = cfg.machine.bw_intersip, cfg.machine.alpha_intersip, "cross-SIP"
    comm_time = 2 * (tp - 1) / tp * bytes_ / bw + 2 * (tp - 1) * alpha
    total_comm_bytes = int(2 * (tp - 1) / tp * bytes_)
    return StageCost(
        name=f"C2 TP AllReduce W_O (TP={tp} ranks, {scope})",
        formula=f"2*(TP-1)/TP * {T_q}*{d}*{b} B / BW + 2(TP-1)*alpha [{scope}]",
        compute_s=0, memory_s=0, comm_s=comm_time,
        bound="comm", visible_s=comm_time,
        flops=0, mem_bytes=0, comm_bytes=total_comm_bytes,
        flops_formula="0",
        mem_formula="0",
        comm_formula=(
            f"PURPOSE: W_O is row-parallel across TP ranks (each rank holds\n"
            f"1/{tp} of W_O's input rows). After the local W_O GEMM, each of\n"
            f"the {tp} TP ranks holds a PARTIAL hidden vector (sum over its\n"
            f"own Q-head slice only). This all-reduce sums those partials so\n"
            f"every rank ends up with the full hidden vector for the next\n"
            f"stage (RMSNorm -> FFN). Fires ONCE per layer over {scope} links.\n"
            f"---\n"
            f"2*(TP-1)/TP * T_q*d*b = 2*({tp}-1)/{tp} * {T_q}*{d}*{b} "
            f"= {total_comm_bytes/1e6:.2f} MB over {scope} at "
            f"{bw/1e9:.0f} GB/s"
        ),
    )


def comm_kv_split_allreduce(cfg: FullConfig) -> StageCost:
    """Extra AllReduce on attention scores when TP > H_kv with head-dim split.

    Ranks sharing a KV head have partial scores; must AllReduce across the
    split_factor group to get final scores per hop.
    """
    if not cfg.kv_replication_needed or cfg.topo.kv_shard_mode != "split":
        return StageCost(
            name="C3 Score AllReduce (head-split)",
            formula="TP <= H_kv or replicate mode: not needed",
            compute_s=0, memory_s=0, comm_s=0, bound="trivial", visible_s=0,
        )
    split = cfg.head_dim_split_factor  # ranks sharing one head
    T_q = cfg.topo.T_q
    S_local = cfg.topo.s_local
    hq_per_pe = cfg.h_q_per_pe
    b = cfg.model.bytes_per_elem
    # score bytes per rank per hop
    bytes_per_hop = hq_per_pe * T_q * S_local * b
    # split-group AllReduce (intra-cube assumed if group fits)
    bw = cfg.machine.bw_intra
    alpha = cfg.machine.alpha_intra
    per_hop = 2 * (split - 1) / split * bytes_per_hop / bw + 2 * (split - 1) * alpha
    total = per_hop * cfg.topo.cp
    total_comm_bytes = int(2 * (split - 1) / split * bytes_per_hop * cfg.topo.cp)
    return StageCost(
        name=f"C3 Score AllReduce ({split}-way, xCP hops)",
        formula=f"2*({split}-1)/{split} * {hq_per_pe}*{T_q}*{S_local}*{b} / BW + latency",
        compute_s=0, memory_s=0, comm_s=total,
        bound="comm", visible_s=total,
        flops=0, mem_bytes=0, comm_bytes=total_comm_bytes,
        flops_formula="0",
        mem_formula="0",
        comm_formula=(
            f"PURPOSE: TP={cfg.topo.tp} > H_kv={cfg.model.h_kv}, so each KV\n"
            f"head is split across {split} TP ranks along the head-dim (d_h\n"
            f"/{split} per rank). Each rank's Q.K^T is therefore only a\n"
            f"PARTIAL dot product; the {split} ranks sharing one KV head\n"
            f"must AllReduce their partial scores to get the true score.\n"
            f"This fires per hop of the ring, so {cfg.topo.cp}x per layer.\n"
            f"To eliminate C3: set KV mode = 'replicate' (costs {split}x KV\n"
            f"memory but no per-hop score AR), or lower TP so TP <= H_kv.\n"
            f"---\n"
            f"CP * 2*(split-1)/split * (H_q/TP)*T_q*S_local*b "
            f"= {cfg.topo.cp} * 2*({split}-1)/{split} * {hq_per_pe}*{T_q}*{S_local}*{b} "
            f"= {total_comm_bytes/1e6:.2f} MB total"
        ),
    )


def all_stages(cfg: FullConfig) -> list[StageCost]:
    """Ordered per-layer attention stages.

    In prefill with CP > 1, C1 (the CP ring) is inserted between S7 and S8
    because it runs concurrently with S5-S7 per hop. Placing it there in the
    stage list makes the per-stage table read left-to-right in the order
    things actually happen. C3 (score AR) sits with C1 for the same reason.
    C2 (TP AllReduce on W_O) sits right after S10.

    In decode, C1's comm is folded into S8, so no separate C1 row.
    """
    stages = [
        stage_rmsnorm(cfg),
        stage_wq(cfg),
        stage_wkv(cfg),
        stage_kv_append(cfg),
        stage_qkT(cfg),
        stage_softmax(cfg),
        stage_pv(cfg),
    ]
    # Concurrent comm (with S5-S7 in prefill) goes here, right after S7.
    if cfg.topo.mode == "prefill" and cfg.topo.cp > 1:
        stages.append(comm_cp_ring(cfg))
    _c3 = comm_kv_split_allreduce(cfg)
    if _c3.visible_s > 0 or _c3.comm_s > 0:
        stages.append(_c3)
    stages.extend([
        stage_merge(cfg),
        stage_normalize(cfg),
        stage_wo(cfg),
    ])
    # C2 fires right after S10 (W_O produces per-TP-rank output; AR combines).
    stages.append(comm_tp_allreduce(cfg))
    return stages


# ── FFN block stages (lightweight; per PE, per layer) ────────────
def stage_ffn_rmsnorm(cfg: FullConfig) -> StageCost:
    T_q = cfg.topo.T_q
    d = cfg.model.hidden
    b = cfg.model.bytes_per_elem
    bytes_ = T_q * d * b * 2
    flops = 4 * T_q * d
    mem_s = bytes_ / cfg.machine.bw_hbm
    return StageCost(
        name="F1 RMSNorm (pre-FFN)",
        formula=f"{T_q}*{d}*{b}*2 = {bytes_} B / BW_HBM",
        compute_s=0, memory_s=mem_s, comm_s=0,
        bound="memory", visible_s=mem_s,
        flops=flops, mem_bytes=bytes_,
        flops_formula=f"4*T_q*d = 4*{T_q}*{d} = {flops}",
        mem_formula=f"2*T_q*d*b = 2*{T_q}*{d}*{b} = {bytes_} B",
    )


def _ffn_gemm(cfg: FullConfig, name: str, ffn_per_pe: int) -> StageCost:
    T_q = cfg.topo.T_q
    d = cfg.model.hidden
    b = cfg.model.bytes_per_elem
    flops = 2 * T_q * d * ffn_per_pe
    weight_B = d * ffn_per_pe * b
    cmp_s, mem_s = _gemm_time(flops, weight_B, cfg)
    vis, bnd = _visible(cmp_s, mem_s, 0)
    return StageCost(
        name=name,
        formula=f"FLOPs = 2*{T_q}*{d}*{ffn_per_pe} = {flops:.2g}; weight = {weight_B/1e6:.1f} MB",
        compute_s=cmp_s, memory_s=mem_s, comm_s=0,
        bound=bnd, visible_s=vis,
        flops=int(flops), mem_bytes=int(weight_B),
        flops_formula=f"2*T_q*d*(ffn/div) = 2*{T_q}*{d}*{ffn_per_pe} = {flops:.3g}",
        mem_formula=f"d*(ffn/div)*b = {d}*{ffn_per_pe}*{b} = {weight_B/1e6:.2f} MB",
    )


def stage_ffn_gate(cfg: FullConfig) -> StageCost:
    ffn_per_pe = max(1, cfg.model.ffn_dim // (cfg.ffn_shard_divisor * max(1, cfg.topo.ep)))
    return _ffn_gemm(cfg, "F2 W_gate GEMM", ffn_per_pe)


def stage_ffn_up(cfg: FullConfig) -> StageCost:
    ffn_per_pe = max(1, cfg.model.ffn_dim // (cfg.ffn_shard_divisor * max(1, cfg.topo.ep)))
    return _ffn_gemm(cfg, "F3 W_up GEMM", ffn_per_pe)


def stage_ffn_swiglu(cfg: FullConfig) -> StageCost:
    """SwiGLU element-wise activation on the intermediate FFN tensor."""
    T_q = cfg.topo.T_q
    ffn_per_pe = max(1, cfg.model.ffn_dim // (cfg.ffn_shard_divisor * max(1, cfg.topo.ep)))
    b = cfg.model.bytes_per_elem
    bytes_ = 3 * T_q * ffn_per_pe * b
    flops = 3 * T_q * ffn_per_pe  # gate * silu(up) approximated as 3 flops/elt
    mem_s = bytes_ / cfg.machine.bw_hbm
    return StageCost(
        name="F4 SwiGLU act",
        formula=f"3*{T_q}*{ffn_per_pe}*{b} = {bytes_} B / BW_HBM",
        compute_s=0, memory_s=mem_s, comm_s=0,
        bound="memory", visible_s=mem_s,
        flops=flops, mem_bytes=bytes_,
        flops_formula=f"~3*T_q*(ffn/div) = 3*{T_q}*{ffn_per_pe} = {flops}",
        mem_formula=f"3*T_q*(ffn/div)*b = 3*{T_q}*{ffn_per_pe}*{b} = {bytes_} B",
    )


def stage_ffn_down(cfg: FullConfig) -> StageCost:
    """W_down: (ffn_per_pe, hidden). Row-parallel over FFN scope."""
    T_q = cfg.topo.T_q
    d = cfg.model.hidden
    b = cfg.model.bytes_per_elem
    ffn_per_pe = max(1, cfg.model.ffn_dim // (cfg.ffn_shard_divisor * max(1, cfg.topo.ep)))
    flops = 2 * T_q * ffn_per_pe * d
    weight_B = ffn_per_pe * d * b
    cmp_s, mem_s = _gemm_time(flops, weight_B, cfg)
    vis, bnd = _visible(cmp_s, mem_s, 0)
    return StageCost(
        name="F5 W_down GEMM",
        formula=f"FLOPs = 2*{T_q}*{ffn_per_pe}*{d} = {flops:.2g}; weight = {weight_B/1e6:.1f} MB",
        compute_s=cmp_s, memory_s=mem_s, comm_s=0,
        bound=bnd, visible_s=vis,
        flops=int(flops), mem_bytes=int(weight_B),
        flops_formula=f"2*T_q*(ffn/div)*d = 2*{T_q}*{ffn_per_pe}*{d} = {flops:.3g}",
        mem_formula=f"(ffn/div)*d*b = {ffn_per_pe}*{d}*{b} = {weight_B/1e6:.2f} MB",
    )


def comm_ffn_allreduce(cfg: FullConfig) -> StageCost:
    """AllReduce on FFN output across the FFN sharding scope."""
    scope = cfg.topo.ffn_shard_scope
    divisor = cfg.ffn_shard_divisor  # TP or TP*CP or TP*CP*DP
    if divisor <= 1:
        return StageCost(
            name="CF1 FFN AllReduce",
            formula="FFN scope = 1: not needed",
            compute_s=0, memory_s=0, comm_s=0,
            bound="trivial", visible_s=0,
        )
    T_q = cfg.topo.T_q
    d = cfg.model.hidden
    b = cfg.model.bytes_per_elem
    bytes_ = T_q * d * b
    # Choose BW/alpha tier based on scope (rough): TP=intra-cube (or inter-cube),
    # +CP=inter-SIP possible, +DP=inter-SIP always.
    if "DP" in scope or "CP" in scope:
        bw = cfg.machine.bw_intersip if cfg.topo.sips_used > 1 else cfg.machine.bw_inter
        alpha = cfg.machine.alpha_intersip if cfg.topo.sips_used > 1 else cfg.machine.alpha_inter
    else:  # TP only
        bw = cfg.machine.bw_intra if cfg.topo.tp_spans_cubes == 1 else cfg.machine.bw_inter
        alpha = cfg.machine.alpha_intra if cfg.topo.tp_spans_cubes == 1 else cfg.machine.alpha_inter
    comm_time = 2 * (divisor - 1) / divisor * bytes_ / bw + 2 * (divisor - 1) * alpha
    total_comm_bytes = int(2 * (divisor - 1) / divisor * bytes_)
    return StageCost(
        name=f"CF1 FFN AllReduce (scope={scope}, x{divisor})",
        formula=f"2*({divisor}-1)/{divisor} * {T_q}*{d}*{b} B / BW + latency",
        compute_s=0, memory_s=0, comm_s=comm_time,
        bound="comm", visible_s=comm_time,
        flops=0, mem_bytes=0, comm_bytes=total_comm_bytes,
        flops_formula="0",
        mem_formula="0",
        comm_formula=(
            f"PURPOSE: W_down is row-parallel across the {divisor} ranks in\n"
            f"the FFN scope (scope={scope}). Each rank produced a PARTIAL\n"
            f"hidden vector after W_down; this all-reduce sums them so every\n"
            f"rank has the full FFN output for the next layer. Larger scope\n"
            f"= less FFN weight memory per PE but bigger AR bill.\n"
            f"---\n"
            f"2*(div-1)/div * T_q*d*b = 2*({divisor}-1)/{divisor} * "
            f"{T_q}*{d}*{b} = {total_comm_bytes/1e6:.2f} MB at "
            f"{bw/1e9:.0f} GB/s (scope={scope})"
        ),
    )


def all_ffn_stages(cfg: FullConfig) -> list[StageCost]:
    return [
        stage_ffn_rmsnorm(cfg),
        stage_ffn_gate(cfg),
        stage_ffn_up(cfg),
        stage_ffn_swiglu(cfg),
        stage_ffn_down(cfg),
        comm_ffn_allreduce(cfg),
    ]
