"""Streamlit app: interactive analytical model for transformer inference
on the SIP architecture.

Run:
    python -m streamlit run tests/analytical_visualization/app.py
"""
from __future__ import annotations

import importlib
import sys

import matplotlib.pyplot as plt
import pandas as pd
import streamlit as st

# Streamlit's hot-reload only re-runs THIS script; imported sub-modules stay
# cached in sys.modules across reruns. When we edit auto_explore.py or
# auto_hardware.py while the Streamlit process is alive, the app would
# keep using the pre-edit versions until a full Ctrl+C + restart. Force
# a reload of our own modules on every rerun so signature changes land
# without a full restart.
for _mod_name in (
    "tests.analytical_visualization.auto_explore",
    "tests.analytical_visualization.auto_hardware",
    "tests.analytical_visualization.autosuggest",
    "tests.analytical_visualization.stage_latencies",
    "tests.analytical_visualization.memory_layout",
    "tests.analytical_visualization.pe_weight_layout",
    "tests.analytical_visualization.tensor_sharding",
    "tests.analytical_visualization.topology_map",
    "tests.analytical_visualization.pipeline_diagram",
    "tests.analytical_visualization.optimization_report",
):
    _m = sys.modules.get(_mod_name)
    if _m is not None:
        importlib.reload(_m)

from tests.analytical_visualization.model_config import (
    FullConfig, ModelConfig, TopologyConfig, MachineParams,
)
from tests.analytical_visualization.model_presets import PRESETS
from tests.analytical_visualization.stage_latencies import all_stages, all_ffn_stages
from tests.analytical_visualization.memory_layout import (
    compute_memory,
    attention_weight_rows,
    ffn_weight_rows,
    kv_cache_rows,
    sum_bytes_all_layers,
    total_weight_bytes_full_model,
    total_kv_bytes_full_model,
)
from tests.analytical_visualization.topology_map import draw_topology
from tests.analytical_visualization.pe_weight_layout import (
    draw_pe_layout, _per_pe_bytes, _q_heads_for_pe, _kv_heads_for_pe,
)
from tests.analytical_visualization.tensor_sharding import draw_tensor_sharding
from tests.analytical_visualization.pipeline_diagram import draw_pipeline
from tests.analytical_visualization.optimization_report import (
    replication_report, optimization_hints,
)
from tests.analytical_visualization.autosuggest import auto_suggest


st.set_page_config(page_title="SIP Attention Analytical Model",
                   layout="wide")

# Tighter spacing so more info fits per page.
st.markdown(
    """
    <style>
      .block-container { padding-top: 1.2rem; padding-bottom: 1rem;
                         padding-left: 1.5rem; padding-right: 1.5rem; }
      div[data-testid="stMetricValue"] { font-size: 1.05rem; }
      div[data-testid="stMetricLabel"] { font-size: 0.75rem; }
      h1 { font-size: 1.6rem !important; margin-bottom: 0.2rem; }
      h2 { font-size: 1.15rem !important; margin-top: 0.4rem; }
      h3 { font-size: 1.0rem  !important; margin-top: 0.3rem; }
      div[data-testid="stExpander"] summary p { font-size: 0.85rem !important; }
      .streamlit-expanderContent { padding-top: 0.3rem !important; }
      section[data-testid="stSidebar"] .stSelectbox label,
      section[data-testid="stSidebar"] .stRadio label,
      section[data-testid="stSidebar"] .stNumberInput label {
          font-size: 0.78rem; }
    </style>
    """,
    unsafe_allow_html=True,
)


def _pick_in(container, label: str, options, default, key: str, help: str = ""):
    """Selectbox helper — returns the selected value from options."""
    return container.selectbox(
        label, options,
        index=options.index(default) if default in options else 0,
        key=key, help=help,
    )


def _pick(label, options, default, key, help=""):
    return _pick_in(st.sidebar, label, options, default, key, help)


# ── Sidebar ─ compact, guided layout ────────────────────────────
st.sidebar.markdown("### SIP config")
st.sidebar.caption(
    "Pick a model, workload, parallelism plan, then tune hardware. "
    "Every change re-runs the analytical model live."
)

with st.sidebar.expander("Model", expanded=True):
    preset_names = list(PRESETS.keys())
    preset_name = st.selectbox(
        "Preset", preset_names,
        index=preset_names.index("Qwen 3 8B"),
        key="preset_select",
        label_visibility="collapsed",
    )
    preset = PRESETS[preset_name]
    st.caption(
        f"{preset.family or '-'} - {preset.attn_type} - "
        f"H_q={preset.model.h_q}, H_kv={preset.model.h_kv}, "
        f"L={preset.model.layers}"
    )
    with st.expander("Override dims", expanded=False):
        d1, d2 = st.columns(2)
        with d1:
            hidden_ = st.number_input("hidden", 128, 65536,
                                       value=preset.model.hidden, step=128)
            h_q_ = st.number_input("H_q", 1, 256, value=preset.model.h_q)
            d_head_ = st.number_input("d_head", 32, 512,
                                       value=preset.model.d_head, step=32)
            layers_ = st.number_input("layers", 1, 200,
                                       value=preset.model.layers)
        with d2:
            ffn_ = st.number_input("ffn_dim", 128, 131072,
                                    value=preset.model.ffn_dim, step=128)
            h_kv_ = st.number_input("H_kv", 1, 256, value=preset.model.h_kv)
            bpe_ = st.selectbox("dtype bytes", [1, 2, 4], index=1)
        model = ModelConfig(
            name=preset.model.name, hidden=hidden_, ffn_dim=ffn_,
            h_q=h_q_, h_kv=h_kv_, d_head=d_head_, layers=layers_,
            bytes_per_elem=bpe_,
        )


with st.sidebar.expander("Workload", expanded=True):
    w1, w2 = st.columns([2, 1])
    with w1:
        s_kv = st.selectbox(
            "Context S_kv",
            [1024, 2048, 4096, 8192, 16384, 32768, 65536,
             131072, 262144, 524288, 1048576],
            index=10, key="s_kv",  # default 1M
        )
    with w2:
        mode = st.selectbox("Mode", ["decode", "prefill"],
                             index=0, key="mode")


# Compute auto-suggest for defaults / apply button
_default_machine = MachineParams()
_default_suggestion = auto_suggest(model, _default_machine, s_kv, mode)

def _snap(v: int, opts: tuple) -> int:
    return min(opts, key=lambda o: abs(o - v))

_CP_OPTS = (1, 2, 4, 8, 16, 32, 48, 64, 96)
_TP_OPTS = (1, 2, 4, 8, 16, 32)
_PP_OPTS = (1, 2, 4, 8, 16, 32)
_DP_OPTS = (1, 2, 4, 8, 16)
_EP_OPTS = (1, 2, 4, 8, 16, 32)

# Auto-reset dropdowns when model preset changes.
if st.session_state.get("_last_preset") != preset_name:
    st.session_state["cp"] = _snap(_default_suggestion.cp, _CP_OPTS)
    st.session_state["tp"] = _snap(_default_suggestion.tp, _TP_OPTS)
    st.session_state["pp"] = _snap(_default_suggestion.pp, _PP_OPTS)
    st.session_state["dp"] = 1
    st.session_state["ep"] = 1
    # Apply the picked cp_placement so the sidebar reflects the packed
    # config (autosuggest may have chosen "pe" packing to minimize cubes).
    st.session_state["cp_placement"] = _default_suggestion.cp_placement
    st.session_state["_last_preset"] = preset_name

with st.sidebar.expander("Parallelism", expanded=True):
    st.caption(
        f"Auto-suggest (memory-min): CP={_default_suggestion.cp}, "
        f"TP={_default_suggestion.tp}, PP={_default_suggestion.pp} "
        f"→ **{_default_suggestion.cubes_used} cube(s)**, "
        f"{_default_suggestion.pes_used} PEs "
        f"(cp_placement={_default_suggestion.cp_placement})"
    )
    # Memory-min apply — snap sliders back to what the caption says.
    if st.button("Apply memory-min", width='stretch',
                  key="_sb_apply_memmin",
                  help="Set sliders to the memory-min config shown above "
                       "(smallest deployment that fits)."):
        st.session_state["cp"] = _snap(_default_suggestion.cp, _CP_OPTS)
        st.session_state["tp"] = _snap(_default_suggestion.tp, _TP_OPTS)
        st.session_state["pp"] = _snap(_default_suggestion.pp, _PP_OPTS)
        st.session_state["dp"] = 1
        st.session_state["cp_placement"] = _default_suggestion.cp_placement
        st.session_state["_pl_active_scope"] = "full"
        st.rerun()
    # Three latency-optimal apply buttons — pick a scope, load the
    # latency-min Pareto config into the sidebar sliders.
    from tests.analytical_visualization.auto_explore import (
        run_auto_explore as _run_auto_explore_sb,
    )
    _sb_scope_meta = {
        "attn": (True,  False, "Attention"),
        "ffn":  (False, True,  "FFN/MoE"),
        "full": (True,  True,  "Attn+FFN"),
    }
    _sb_cols = st.columns(3)
    _sb_clicks = {}
    for _sb_scope, _sb_col in zip(("attn", "ffn", "full"), _sb_cols):
        with _sb_col:
            _sb_clicks[_sb_scope] = st.button(
                _sb_scope_meta[_sb_scope][2], width='stretch',
                key=f"_sb_apply_{_sb_scope}",
                help=f"Apply latency-optimal parallelism for "
                     f"{_sb_scope_meta[_sb_scope][2]} scope.",
            )
    for _sb_scope, _sb_click in _sb_clicks.items():
        if not _sb_click:
            continue
        _sb_ia, _sb_ff, _sb_lbl = _sb_scope_meta[_sb_scope]
        with st.spinner(f"Finding latency-optimal parallelism ({_sb_lbl})..."):
            _sb_res = _run_auto_explore_sb(
                model, _default_machine, s_kv=s_kv, mode=mode,
                include_attention=_sb_ia, include_ffn=_sb_ff,
            )
        if not _sb_res.pareto_scores:
            st.error(
                f"No feasible parallelism for {_sb_lbl}. Try raising "
                "`pe_hbm_gb` or reducing `s_kv`."
            )
        else:
            _sb_best = min(_sb_res.pareto_scores,
                            key=lambda s: s.total_latency_ns)
            st.session_state["cp"] = _snap(_sb_best.cp, _CP_OPTS)
            st.session_state["tp"] = _snap(_sb_best.tp, _TP_OPTS)
            st.session_state["pp"] = _snap(_sb_best.pp, _PP_OPTS)
            st.session_state["dp"] = _snap(_sb_best.dp, _DP_OPTS)
            st.session_state["tp_placement"] = _sb_best.tp_placement
            st.session_state["cp_placement"] = _sb_best.cp_placement
            st.session_state["cp_ring_variant"] = _sb_best.cp_ring_variant
            st.session_state["kv_mode"] = _sb_best.kv_shard_mode
            _sb_tp2 = _snap(_sb_best.tp, _TP_OPTS)
            _sb_cp2 = _snap(_sb_best.cp, _CP_OPTS)
            _sb_dp2 = _snap(_sb_best.dp, _DP_OPTS)
            _sb_ffn_map = {
                "TP":       f"TP only (div={max(1, _sb_tp2)})",
                "TP+CP":    f"TP*CP (div={max(1, _sb_tp2 * _sb_cp2)})",
                "TP+CP+DP": f"TP*CP*DP (div={max(1, _sb_tp2 * _sb_cp2 * _sb_dp2)})",
            }
            st.session_state["ffn_scope_label"] = \
                _sb_ffn_map[_sb_best.ffn_shard_scope]
            # Track for the Physical Layout scope filter, so the tables
            # below also filter to this scope.
            st.session_state["_pl_active_scope"] = _sb_scope
            st.rerun()

    # Sharding dims — arrange in 2x3 grid to save vertical space
    p_row1 = st.columns(3)
    with p_row1[0]:
        cp = st.selectbox(
            "CP", list(_CP_OPTS),
            index=_CP_OPTS.index(st.session_state.get(
                "cp", _snap(_default_suggestion.cp, _CP_OPTS))),
            key="cp",
        )
    with p_row1[1]:
        tp = st.selectbox(
            "TP", list(_TP_OPTS),
            index=_TP_OPTS.index(st.session_state.get(
                "tp", _snap(_default_suggestion.tp, _TP_OPTS))),
            key="tp",
        )
    with p_row1[2]:
        pp = st.selectbox(
            "PP", list(_PP_OPTS),
            index=_PP_OPTS.index(st.session_state.get(
                "pp", _snap(_default_suggestion.pp, _PP_OPTS))),
            key="pp",
        )
    p_row2 = st.columns(2)
    with p_row2[0]:
        dp = st.selectbox(
            "DP", list(_DP_OPTS),
            index=_DP_OPTS.index(st.session_state.get("dp", 1)),
            key="dp",
        )
    with p_row2[1]:
        ep = st.selectbox(
            "EP (MoE)", list(_EP_OPTS),
            index=_EP_OPTS.index(st.session_state.get("ep", 1)),
            key="ep",
        )

    st.markdown("**Placement (PE-level vs cube-level)**")
    pl_col1, pl_col2 = st.columns(2)
    with pl_col1:
        tp_placement = st.radio(
            "TP on", options=["pe", "cube"], index=0,
            key="tp_placement", horizontal=True,
            help=("pe: TP shares PEs of one cube (intra-cube AllReduce). "
                  "cube: each TP rank owns a whole cube (inter-cube AllReduce)."),
        )
    with pl_col2:
        cp_placement = st.radio(
            "CP on", options=["pe", "cube"], index=1,
            key="cp_placement", horizontal=True,
            help=("pe: CP ring runs among PEs of one cube (intra-cube). "
                  "cube: each CP rank owns cubes (inter-cube ring)."),
        )

    st.markdown("**Sharding modes**")
    # FFN scope labels show the effective divisor (dynamic per CP/TP/DP).
    _ffn_labels = {
        "TP": f"TP only (div={max(1,tp)})",
        "TP+CP": f"TP*CP (div={max(1,tp*cp)})",
        "TP+CP+DP": f"TP*CP*DP (div={max(1,tp*cp*dp)})",
    }
    _ffn_choice = st.radio(
        "FFN shard scope",
        options=list(_ffn_labels.values()),
        index=1, key="ffn_scope_label", horizontal=True,
        help=("Divisor = number of ranks the FFN weights are split across. "
              "Larger divisor = less memory/PE, more AllReduce cost."),
    )
    # Map label -> raw scope key.
    ffn_shard_scope = {v: k for k, v in _ffn_labels.items()}[_ffn_choice]

    sip_topology = st.radio(
        "SIP interconnect",
        options=["ring", "mesh2d", "torus2d"],
        index=0, key="sip_topo", horizontal=True,
        help=("ring: 1D chain + wrap. mesh2d: 2D grid, no wrap. "
              "torus2d: 2D grid + wrap."),
    )
    cp_ring_variant = st.radio(
        "CP ring: what rotates",
        options=["kv", "qoml"], index=0, key="cp_ring_variant",
        horizontal=True,
        help=("kv: K,V shards rotate each hop (Q,O,m,l stay local). "
              "Bytes/hop scales with S_local*H_kv - good for prefill.\n"
              "qoml: Q + running (O,m,l) rotate (K,V stay local). "
              "Bytes/hop scales with T_q*H_q - much cheaper for decode."),
    )
    kv_shard_mode = "split"
    if tp > model.h_kv:
        _rep_factor = max(1, tp // model.h_kv)
        kv_shard_mode = st.radio(
            f"KV mode (TP={tp} > H_kv={model.h_kv})",
            options=["split", "replicate"], index=0,
            key="kv_mode", horizontal=True,
            help=(f"split: fractional heads (KV/PE shrinks by {tp/max(1,model.h_kv):.1f}x, "
                  f"adds Score AllReduce across {_rep_factor} ranks). "
                  f"replicate: whole heads per PE, duplicated {_rep_factor}x, "
                  f"no Score AllReduce."),
        )


with st.sidebar.expander("Hardware", expanded=False):
    hw_tab_pe, hw_tab_net = st.tabs(["Per-PE", "Interconnect"])
    with hw_tab_pe:
        hp1, hp2 = st.columns(2)
        with hp1:
            pe_hbm_gb = st.selectbox(
                "PE HBM (GB)", [3.0, 6.0, 12.0, 24.0, 48.0, 96.0],
                index=1, key="pe_hbm",
            )
            bw_hbm = st.selectbox(
                "HBM BW GB/s", [128.0, 256.0, 512.0, 1024.0, 2048.0],
                index=1, key="bw_hbm",
            )
        with hp2:
            peak_tflops = st.selectbox(
                "TFLOPs/PE", [2.0, 4.0, 8.0, 16.0, 32.0, 64.0],
                index=2, key="tflops",
            )
            compute_util = st.selectbox(
                "Compute util", [0.3, 0.5, 0.7, 0.8, 0.9, 1.0],
                index=3, key="util",
            )
    with hw_tab_net:
        n1, n2 = st.columns(2)
        with n1:
            bw_intra = st.selectbox(
                "intra-cube GB/s",
                [128.0, 256.0, 512.0, 1024.0, 2048.0, 4096.0],
                index=2, key="bw_intra",
            )
            bw_inter = st.selectbox(
                "inter-cube GB/s",
                [32.0, 64.0, 128.0, 256.0, 512.0, 900.0],
                index=2, key="bw_inter",
            )
            bw_intersip = st.selectbox(
                "inter-SIP GB/s",
                [12.5, 25.0, 50.0, 100.0, 200.0, 400.0],
                index=2, key="bw_intersip",
            )
        with n2:
            alpha_intra = st.selectbox(
                "a_intra ns", [5.0, 10.0, 20.0, 50.0, 100.0],
                index=2, key="a_intra",
            )
            alpha_inter = st.selectbox(
                "a_inter ns", [50.0, 100.0, 200.0, 500.0, 1000.0],
                index=1, key="a_inter",
            )
            alpha_intersip = st.selectbox(
                "a_intersip ns",
                [500.0, 1000.0, 2000.0, 5000.0, 10000.0],
                index=1, key="a_intersip",
            )

machine = MachineParams(
    pe_hbm_gb=pe_hbm_gb, peak_tflops_f16=peak_tflops,
    bw_hbm_gbs=bw_hbm, bw_intra_gbs=bw_intra,
    bw_inter_gbs=bw_inter, bw_intersip_gbs=bw_intersip,
    alpha_intra_ns=alpha_intra, alpha_inter_ns=alpha_inter,
    alpha_intersip_ns=alpha_intersip,
    compute_util=compute_util,
)

topo = TopologyConfig(cp=cp, tp=tp, pp=pp, dp=dp, ep=ep,
                       s_kv=s_kv, mode=mode,
                       kv_shard_mode=kv_shard_mode,
                       ffn_shard_scope=ffn_shard_scope,
                       sip_topology=sip_topology,
                       tp_placement=tp_placement,
                       cp_placement=cp_placement,
                       cp_ring_variant=cp_ring_variant)
cfg = FullConfig(model=model, topo=topo, machine=machine)


# ── Main pane ─────────────────────────────────────────────────────
st.title(f"SIP Analytical Model — {model.name}")

st.caption(
    f"Family: **{preset.family or '-'}**  |  Attn: **{preset.attn_type}**  |  "
    f"H_q={model.h_q}, H_kv={model.h_kv}, hidden={model.hidden}, "
    f"ffn={model.ffn_dim}, layers={model.layers}"
)
if preset.note:
    st.info(f"Note: {preset.note}")


# Recompute auto-suggest with the *current* machine (in case user changed HBM etc.)
suggestion = auto_suggest(model, machine, s_kv, mode)

# Compact top summary: auto-suggest + current in one row of tight bullets.
_sug_status = "OK" if suggestion.fits else "NO FIT"
_sug_line = (
    f"**Auto-suggest:** CP={suggestion.cp} - TP={suggestion.tp} - "
    f"PP={suggestion.pp} - PEs={suggestion.pes_used} - SIPs={suggestion.sips_used} "
    f"({_sug_status}, slack {suggestion.slack_gb:.2f} GB / "
    f"{machine.pe_hbm_gb:.1f} GB budget)"
)
_cur_line = (
    f"**Current:** CP={cp}@{cp_placement} - TP={tp}@{tp_placement} - "
    f"PP={pp} - DP={dp} - EP={ep} - PEs={topo.total_pes} - "
    f"Cubes={topo.cubes_used} ({topo.pes_per_cube_used}/cube live) - "
    f"SIPs={topo.sips_used} - FFN={cfg.topo.ffn_shard_scope} - "
    f"SIPnet={cfg.topo.sip_topology}"
)
sum_l, sum_r = st.columns([1, 1])
with sum_l:
    st.markdown(_sug_line)
with sum_r:
    st.markdown(_cur_line)

# Consolidated warnings (only if any apply)
_warnings = []
if not topo.placement_valid:
    _spill_cubes = (topo.intra_cube_dims + topo.pes_per_cube_hw - 1) // topo.pes_per_cube_hw
    _warnings.append(
        f"Intra-cube demand = {topo.intra_cube_dims} PEs > "
        f"{topo.pes_per_cube_hw} PEs/cube -> spills to {_spill_cubes} cubes "
        f"per group (AllReduce/ring drops from intra-cube "
        f"{machine.bw_intra_gbs:.0f} to inter-cube "
        f"{machine.bw_inter_gbs:.0f} GB/s)."
    )
if tp > model.h_kv:
    # Compute per-PE KV bytes under each mode so the difference is visible.
    from tests.analytical_visualization.memory_layout import per_pe_kv_cache_bytes
    import copy as _copy
    _cfg_split = FullConfig(model=model, topo=TopologyConfig(
        **{**topo.__dict__, "kv_shard_mode": "split"}), machine=machine)
    _cfg_rep = FullConfig(model=model, topo=TopologyConfig(
        **{**topo.__dict__, "kv_shard_mode": "replicate"}), machine=machine)
    _kv_split_gb = per_pe_kv_cache_bytes(_cfg_split) / 1e9
    _kv_rep_gb = per_pe_kv_cache_bytes(_cfg_rep) / 1e9
    _kv_now = _kv_split_gb if kv_shard_mode == "split" else _kv_rep_gb
    _warnings.append(
        f"KV mode: **{kv_shard_mode}** (TP={tp} > H_kv={model.h_kv}) - "
        f"per-PE KV: split={_kv_split_gb:.2f} GB vs "
        f"replicate={_kv_rep_gb:.2f} GB (currently {_kv_now:.2f} GB). "
        f"Split adds one Score AllReduce over "
        f"{max(1, tp // model.h_kv)} ranks per hop."
    )
if cfg.kv_replication_needed:
    _warnings.append(f"Head-dim split: {cfg.head_dim_split_factor} ranks share one KV head")
if topo.tp_spans_cubes > 1:
    _warnings.append(f"TP spans {topo.tp_spans_cubes} cubes "
                     f"(AllReduce goes cross-cube)")
if topo.cp_inter_sip_hops > 0:
    _warnings.append(f"CP ring crosses {topo.cp_inter_sip_hops} SIP boundary(ies) "
                     f"at {machine.bw_intersip_gbs:.0f} GB/s")
if _warnings:
    st.warning(" - ".join(_warnings))


# ── Tabs ─────────────────────────────────────────────────────────
# One tab per auto feature; each has TWO sweep buttons (Attention only /
# Attn + FFN/MoE) so the user can toggle scope without switching tabs.
# Auto Suggest is renamed "Parallelism" since it only varies parallelism
# knobs (hardware is held fixed at the sidebar values).
tab_layout, tab_auto, tab_memory, tab_stages, tab_compare, tab_hw = st.tabs([
    "Physical layout",
    "Auto Suggest Parallelism",
    "Memory breakdown", "Per-stage latency",
    "Save & compare", "Auto Hardware",
])


# ── TAB 1: Physical layout ──────────────────────────────────────
with tab_layout:
    # Quick shortcut: three buttons that pick the latency-optimal Pareto
    # config for a chosen scope and load it into the sidebar. Layout
    # re-renders immediately on the next Streamlit run with the new config.
    from tests.analytical_visualization.auto_explore import (
        run_auto_explore as _run_auto_explore_layout,
    )

    st.markdown(
        "**Quick auto-suggest layout** — pick a scope; the latency-optimal "
        "Pareto config for that scope is loaded into the sidebar and this "
        "layout redraws around it."
    )
    _pl1, _pl2, _pl3, _pl4 = st.columns([1, 1, 1, 1])
    with _pl1:
        _pl_run_attn = st.button("Run sweep — Attention",
                                   type="primary", width='stretch',
                                   key="_pl_run_attn")
    with _pl2:
        _pl_run_ffn = st.button("Run sweep — FFN/MoE",
                                  type="primary", width='stretch',
                                  key="_pl_run_ffn")
    with _pl3:
        _pl_run_full = st.button("Run sweep — Attn + FFN/MoE",
                                   type="primary", width='stretch',
                                   key="_pl_run_full")

    _pl_scope_meta = {
        "attn": (True,  False, "Attention only"),
        "ffn":  (False, True,  "FFN / MoE only"),
        "full": (True,  True,  "Attn + FFN/MoE"),
    }
    _pl_button_click = {
        "attn": _pl_run_attn, "ffn": _pl_run_ffn, "full": _pl_run_full,
    }
    for _pl_scope, _pl_clicked in _pl_button_click.items():
        if not _pl_clicked:
            continue
        _pl_ia, _pl_ff, _pl_lbl = _pl_scope_meta[_pl_scope]
        with st.spinner(f"Finding latency-optimal parallelism ({_pl_lbl})..."):
            _pl_res = _run_auto_explore_layout(
                model, machine, s_kv=s_kv, mode=mode,
                include_attention=_pl_ia, include_ffn=_pl_ff,
            )
        if not _pl_res.pareto_scores:
            st.error(
                f"No feasible parallelism for {_pl_lbl} — every config "
                "exceeds per-PE HBM. Try raising `pe_hbm_gb` or reducing "
                "`s_kv`."
            )
        else:
            _pl_best = min(_pl_res.pareto_scores,
                            key=lambda s: s.total_latency_ns)
            # Load the winning config into the sidebar's session state.
            st.session_state["cp"] = _pl_best.cp
            st.session_state["tp"] = _pl_best.tp
            st.session_state["pp"] = _pl_best.pp
            st.session_state["dp"] = _pl_best.dp
            st.session_state["tp_placement"] = _pl_best.tp_placement
            st.session_state["cp_placement"] = _pl_best.cp_placement
            st.session_state["cp_ring_variant"] = _pl_best.cp_ring_variant
            st.session_state["kv_mode"] = _pl_best.kv_shard_mode
            _pl_tp, _pl_cp, _pl_dp = _pl_best.tp, _pl_best.cp, _pl_best.dp
            _pl_ffn_map = {
                "TP":       f"TP only (div={max(1, _pl_tp)})",
                "TP+CP":    f"TP*CP (div={max(1, _pl_tp * _pl_cp)})",
                "TP+CP+DP": f"TP*CP*DP (div={max(1, _pl_tp * _pl_cp * _pl_dp)})",
            }
            st.session_state["ffn_scope_label"] = \
                _pl_ffn_map[_pl_best.ffn_shard_scope]
            # Persist the chosen scope so subsequent layout renders filter
            # stage tables (attn/ffn/full).
            st.session_state["_pl_active_scope"] = _pl_scope
            st.success(
                f"Loaded latency-optimal {_pl_lbl} config: CP={_pl_best.cp} "
                f"TP={_pl_best.tp} PP={_pl_best.pp} DP={_pl_best.dp} → "
                f"{_pl_best.latency_ms:.2f} ms, {_pl_best.pes_used} PEs. "
                "Redrawing layout..."
            )
            st.rerun()

    # Determine the effective scope for filtering the tables below. Default
    # to "full" so a fresh page load / sidebar-driven config shows everything.
    _pl_active_scope = st.session_state.get("_pl_active_scope", "full")
    _pl_show_attn = _pl_active_scope in ("attn", "full")
    _pl_show_ffn = _pl_active_scope in ("ffn", "full")
    _pl_active_label = _pl_scope_meta[_pl_active_scope][2]
    st.markdown(f"### Layout scope: **{_pl_active_label}**")

    st.divider()

    # Row A: pipeline diagram (color-coded) + per-stage table (formulas).
    st.markdown("**Per-layer pipeline (color = dominant bound)**")
    fig_pipe = draw_pipeline(cfg)
    st.pyplot(fig_pipe, width='stretch')
    plt.close(fig_pipe)

    with st.expander("Per-stage latency (attention + FFN) with formulas",
                      expanded=True):
        import html as _html
        _stages_attn = all_stages(cfg)
        _stages_ffn = all_ffn_stages(cfg)
        _bound_color = {
            "compute": "#d6eaf8",
            "memory":  "#fdebd0",
            "comm":    "#f5b7b1",
            "trivial": "#eaeded",
        }

        def _vsplit(formula: str) -> str:
            """Break 'X = Y = Z' into vertical `= X` / `= Y` / `= Z` lines.
            If the string contains a 'PURPOSE:' preamble ending with a line
            of '---', keep that preamble verbatim and only vsplit the tail.
            """
            if not formula or formula in ("-", "0"):
                return formula
            _preamble = ""
            _body = formula
            if "PURPOSE:" in formula and "\n---\n" in formula:
                _preamble, _body = formula.split("\n---\n", 1)
                _preamble = _preamble.rstrip() + "\n---"
            pieces = [p.strip() for p in _body.split(" = ")]
            _vsplit_body = "\n".join(f"= {p}" for p in pieces)
            return f"{_preamble}\n{_vsplit_body}" if _preamble else _vsplit_body

        # Prefill: per-hop ring is concurrent with S5/S6/S7 compute -
        # annotate them. In decode the O/m/l all-reduce is already
        # folded into the S8 row, so no separate annotation needed.
        _per_hop_ring_us = 0.0
        _ring_variant_desc = ""
        if topo.mode == "prefill" and topo.cp > 1:
            _c1 = next((_x for _x in _stages_attn
                         if _x.name.startswith("C1 ")), None)
            if _c1 is not None and _c1.comm_s > 0:
                _hops = max(1, topo.cp - 1)
                _per_hop_ring_us = (_c1.comm_s * 1e6) / _hops
                _ring_variant_desc = ("K/V" if topo.cp_ring_variant == "kv"
                                       else "Q+O/m/l")

        def _fmt_time(sec: float) -> str:
            """Auto-scale time: ns for < 1 us, us for < 10 ms, else ms."""
            if sec <= 0:
                return "0"
            ns = sec * 1e9
            if ns < 1000:
                return f"{ns:.1f} ns"
            us = sec * 1e6
            if us < 1000:
                return f"{us:.2f} us"
            ms = sec * 1e3
            return f"{ms:.2f} ms"

        def _combined_formula(s):
            parts = []
            if s.compute_s > 0 or (s.flops and s.flops > 0):
                parts.append(
                    f"Compute:\n{_vsplit(s.flops_formula)}\n"
                    f"-> {_fmt_time(s.compute_s)}"
                )
            if s.memory_s > 0 or (s.mem_bytes and s.mem_bytes > 0):
                parts.append(
                    f"Memory:\n{_vsplit(s.mem_formula)}\n"
                    f"-> {_fmt_time(s.memory_s)}"
                )
            if s.comm_s > 0 or (s.comm_bytes and s.comm_bytes > 0):
                # If this is a concurrent-comm stage (e.g. C1 CP ring in
                # prefill) that gets overlapped with compute, visible < raw.
                # Show both so the '0' in the Time column makes sense.
                _tail = f"-> raw comm: {_fmt_time(s.comm_s)}"
                if s.visible_s < s.comm_s * 0.999:
                    _tail += (
                        f"\n-> visible: {_fmt_time(s.visible_s)}  "
                        f"(concurrent S5-S7 compute hides "
                        f"{_fmt_time(s.comm_s - s.visible_s)} of it)"
                    )
                parts.append(
                    f"Comm:\n{_vsplit(s.comm_formula)}\n{_tail}"
                )
            # Prefill: S5/S6/S7 run concurrently with the CP ring.
            if (_per_hop_ring_us > 0
                    and s.name.startswith(("S5 ", "S6 ", "S7 "))):
                # convert us to seconds for _fmt_time
                _per_hop_s = _per_hop_ring_us / 1e6
                parts.append(
                    f"Concurrent CP ring ({_ring_variant_desc}, per hop):\n"
                    f"~{_fmt_time(_per_hop_s)} in flight while this compute runs\n"
                    f"(aggregated in C1 row; overlap kept only excess vs compute)"
                )
            return "\n\n".join(parts) if parts else "-"

        # Compact CSS: rows auto-size to text; tight padding; columns fit content.
        _stage_table_css = """
        <style>
        table.stage-tbl { border-collapse: collapse; width: auto;
                           font-size: 11.5px; margin-bottom: 6px; }
        table.stage-tbl th { text-align: left; padding: 3px 8px;
                             background: #e9ecef; border: 1px solid #ced4da;
                             font-weight: 600; }
        table.stage-tbl td { padding: 3px 8px; border: 1px solid #dee2e6;
                             vertical-align: top; }
        table.stage-tbl td.stg { white-space: nowrap; font-weight: 600; }
        table.stage-tbl td.bnd { text-align: center; font-weight: 600;
                                 text-transform: uppercase; font-size: 10px;
                                 white-space: nowrap; }
        table.stage-tbl td.us  { text-align: right; font-family: Consolas, monospace;
                                 white-space: nowrap; }
        table.stage-tbl pre { margin: 0; font-family: Consolas, monospace;
                              font-size: 11px; line-height: 1.15;
                              white-space: pre; }
        </style>
        """

        def _render_table(stages, title):
            rows_html = []
            for s in stages:
                bg = _bound_color.get(s.bound, "#eaeded")
                formula = _html.escape(_combined_formula(s))
                rows_html.append(
                    f'<tr style="background:{bg};">'
                    f'<td class="stg">{_html.escape(s.name)}</td>'
                    f'<td><pre>{formula}</pre></td>'
                    f'<td class="bnd">{s.bound}</td>'
                    f'<td class="us">{_fmt_time(s.visible_s)}</td>'
                    f'</tr>'
                )
            html = (
                _stage_table_css
                + f'<b>{_html.escape(title)}</b>'
                + '<table class="stage-tbl">'
                + '<thead><tr><th>Stage</th><th>Formula</th>'
                  '<th>Bound</th><th>Time</th></tr></thead>'
                + '<tbody>' + ''.join(rows_html) + '</tbody></table>'
            )
            # Prefer st.html (reliable) with a markdown fallback.
            if hasattr(st, "html"):
                st.html(html)
            else:
                st.markdown(html, unsafe_allow_html=True)
            _t = sum(s.visible_s for s in stages)
            st.caption(
                f"{title} total: **{_fmt_time(_t)}**. "
                f"Compute time = FLOPs / ({cfg.machine.peak_tflops_f16:.1f} TFLOPs "
                f"x {cfg.machine.compute_util:.0%} util). "
                f"Memory time = bytes / {cfg.machine.bw_hbm_gbs:.0f} GB/s HBM. "
                f"Bound = max(compute, memory, comm) - row color: "
                f"blue=compute, orange=memory, red=comm, grey=trivial."
            )

        if _pl_show_attn:
            _render_table(_stages_attn, "Attention")
        if _pl_show_ffn:
            _render_table(_stages_ffn, "FFN")

        with st.expander("Symbol glossary (what does T_q, d_h, div... mean?)",
                          expanded=False):
            _glossary = [
                ("Model dims", [
                    ("d",       f"hidden dim (current: {model.hidden})"),
                    ("d_h",     f"per-head dim (current: {model.d_head})"),
                    ("H_q",     f"total query heads (current: {model.h_q})"),
                    ("H_kv",    f"total KV heads (current: {model.h_kv})"),
                    ("ffn",     f"FFN intermediate dim (current: {model.ffn_dim})"),
                    ("L",       f"transformer layers (current: {model.layers})"),
                    ("b",       f"bytes per element (dtype size) "
                                f"(current: {model.bytes_per_elem} -> "
                                f"{'fp8' if model.bytes_per_elem == 1 else 'bf16/fp16' if model.bytes_per_elem == 2 else 'fp32'})"),
                ]),
                ("Workload", [
                    ("S_kv",    f"global KV cache length / context (current: {topo.s_kv:,})"),
                    ("S_local", f"per-CP-rank KV length = S_kv / CP "
                                f"(current: {topo.s_local:,})"),
                    ("T_q",     f"query tokens processed this step "
                                f"(decode=1, prefill=S_local; current: {topo.T_q})"),
                    ("mode",    f"decode | prefill (current: {topo.mode})"),
                ]),
                ("Parallelism (degrees)", [
                    ("TP",      f"tensor parallelism (attn heads / FFN cols) "
                                f"(current: {topo.tp}, on {topo.tp_placement})"),
                    ("CP",      f"context parallelism (sequence axis) "
                                f"(current: {topo.cp}, on {topo.cp_placement})"),
                    ("PP",      f"pipeline parallelism (layer stages) (current: {topo.pp})"),
                    ("DP",      f"data parallelism (model replicas) (current: {topo.dp})"),
                    ("EP",      f"expert parallelism (MoE) (current: {topo.ep})"),
                    ("div",     f"FFN shard divisor = product of dims in "
                                f"'FFN shard scope' (current: {cfg.ffn_shard_divisor})"),
                    ("split",   f"# ranks sharing one KV head via head-dim split "
                                f"(TP/H_kv when TP > H_kv; current: "
                                f"{cfg.head_dim_split_factor if cfg.kv_replication_needed else 1})"),
                ]),
                ("Hardware / machine", [
                    ("BW_HBM",         f"per-PE HBM bandwidth "
                                        f"(current: {machine.bw_hbm_gbs:.0f} GB/s)"),
                    ("BW_intra",       f"intra-cube PE<->PE bandwidth "
                                        f"(current: {machine.bw_intra_gbs:.0f} GB/s)"),
                    ("BW_inter",       f"inter-cube (D2D) bandwidth "
                                        f"(current: {machine.bw_inter_gbs:.0f} GB/s)"),
                    ("BW_intersip",    f"inter-SIP (C2C) bandwidth "
                                        f"(current: {machine.bw_intersip_gbs:.0f} GB/s)"),
                    ("alpha_intra",    f"per-hop latency intra-cube "
                                        f"(current: {machine.alpha_intra_ns:.0f} ns)"),
                    ("alpha_inter",    f"per-hop latency inter-cube "
                                        f"(current: {machine.alpha_inter_ns:.0f} ns)"),
                    ("alpha_intersip", f"per-hop latency inter-SIP "
                                        f"(current: {machine.alpha_intersip_ns:.0f} ns)"),
                    ("peak TFLOPs",    f"per-PE peak compute "
                                        f"(current: {machine.peak_tflops_f16:.1f} TFLOPs f16)"),
                    ("util",           f"achievable compute utilization "
                                        f"(current: {machine.compute_util:.0%})"),
                ]),
                ("Stage prefixes", [
                    ("S1 .. S10", "attention forward stages (RMSNorm, W_Q, W_K+W_V, "
                                   "KV append, Q.K^T, softmax, P.V, merge, "
                                   "normalize, W_O)"),
                    ("F1 .. F5",  "FFN forward stages (RMSNorm, W_gate, W_up, "
                                   "SwiGLU, W_down)"),
                    ("C1",         "CP K/V ring (K,V passed hop-by-hop)"),
                    ("C2",         "TP AllReduce on W_O output"),
                    ("C3",         "Score AllReduce when TP > H_kv with head-dim split"),
                    ("CF1",        "FFN AllReduce over the FFN shard scope"),
                ]),
                ("Notation in formulas", [
                    ("H_q / TP",    "Q heads per PE (or per TP rank)"),
                    ("H_kv / TP",   "KV heads per PE (fractional allowed in split mode)"),
                    ("ffn / div",   "FFN cols/rows per PE after sharding"),
                    ("2 * A * B * C", "GEMM FLOPs = 2 * M * K * N (multiply + add per fma)"),
                    ("bytes * b",   "byte count (b = bytes/element)"),
                    ("-> X us",     "resulting latency after dividing by BW or peak"),
                ]),
            ]
            for section, items in _glossary:
                st.markdown(f"**{section}**")
                _rows = "".join(
                    f'<tr><td style="padding:2px 8px;font-family:Consolas,monospace;'
                    f'white-space:nowrap;color:#0d47a1;font-weight:600;">{sym}</td>'
                    f'<td style="padding:2px 8px;">{desc}</td></tr>'
                    for sym, desc in items
                )
                st.html(
                    f'<table style="font-size:11.5px;border-collapse:collapse;'
                    f'margin-bottom:6px;">{_rows}</table>'
                )

    st.divider()

    # Row B: topology diagram (full width - it can be tall for multi-SIP)
    with st.expander(
        f"Physical layout - {topo.total_pes} PEs, {topo.cubes_used} cubes, "
        f"{topo.sips_used} SIP(s), inter-SIP: {topo.sip_topology}",
        expanded=True,
    ):
        fig_topo = draw_topology(cfg)
        st.pyplot(fig_topo, width='stretch')
        plt.close(fig_topo)
        l1, l2, l3, l4 = st.columns(4)
        l1.metric("Cubes", topo.cubes_used)
        l2.metric("SIPs", topo.sips_used)
        l3.metric("CP intra-SIP hops", topo.cp_intra_sip_hops)
        l4.metric("CP inter-SIP hops", topo.cp_inter_sip_hops)

    st.divider()

    # Row C: two-column - LEFT per-PE memory (pie + summary), RIGHT tensor sharding
    mem_layout = compute_memory(cfg)
    col_mem, col_shard = st.columns([1, 1], gap="medium")

    with col_mem:
        st.markdown("**Per-PE memory footprint**")
        labels = ["Weights", "KV cache", "Transient", "Slack"]
        vals = [mem_layout.weights_bytes, mem_layout.kv_cache_bytes,
                mem_layout.transient_bytes, mem_layout.slack_bytes]
        colors = ["#3a86ff", "#8338ec", "#ffbe0b", "#adb5bd"]
        fig_mem_l, ax_mem_l = plt.subplots(figsize=(4.2, 4.2))
        ax_mem_l.pie(vals,
                      labels=[f"{l}\n{v/1e9:.2f} GB"
                              for l, v in zip(labels, vals)],
                      colors=colors, autopct="%1.1f%%", startangle=90,
                      textprops={"fontsize": 8})
        ax_mem_l.set_title(f"Per-PE (budget: "
                            f"{mem_layout.budget_bytes/1e9:.1f} GB)",
                            fontsize=10)
        st.pyplot(fig_mem_l, width='stretch')
        plt.close(fig_mem_l)
        if mem_layout.over_budget:
            st.error(f"OVER BUDGET by "
                     f"{(mem_layout.used_bytes - mem_layout.budget_bytes)/1e9:.2f} GB")
        else:
            st.caption(
                f"Weights {mem_layout.weights_bytes/1e9:.2f} GB - "
                f"KV {mem_layout.kv_cache_bytes/1e9:.2f} GB "
                f"(S_local={topo.s_local:,}) - "
                f"Transient {mem_layout.transient_bytes/1e9:.2f} GB - "
                f"Slack {mem_layout.slack_bytes/1e9:.2f} GB"
            )

    with col_shard:
        st.markdown("**Tensor sharding (which dim is split)**")
        st.caption(
            "W_Q/W_K/W_V/W_gate/W_up: **output cols**. "
            "W_O/W_down: **input rows**. "
            "KV cache: **rows by CP** x **cols by TP**."
        )
        s_c1, s_c2 = st.columns([1, 2])
        with s_c1:
            my_pe_view = st.number_input("PE index",
                                           min_value=0,
                                           max_value=max(0, topo.tp - 1),
                                           value=0, key="my_pe_view")
        with s_c2:
            show_physical = st.checkbox(
                "Enlarge + show physical PE/cube mapping",
                value=False, key="tsv_physical",
                help=("Larger figure with each shard cell labelled by owning "
                      "cube/PE. Especially useful for KV cache when CP is "
                      "split across cubes and PEs."),
            )
        fig_shard = draw_tensor_sharding(cfg, my_pe=my_pe_view,
                                          show_physical=show_physical)
        st.pyplot(fig_shard, width='stretch')
        plt.close(fig_shard)

    st.divider()

    # Row D: PE-level view (wide, full-width) inside an expander
    with st.expander(
        f"PE-level view of one CP group (TP={topo.tp}, "
        f"H_q={model.h_q}, H_kv={model.h_kv})",
        expanded=False,
    ):
        st.caption(
            f"One CP rank's {topo.tp} PE(s). Attention weights replicated "
            f"across CP groups; KV sharded by CP (S/{topo.cp} tokens) x TP heads; "
            f"FFN by TP x EP."
        )
        fig_pe = draw_pe_layout(cfg)
        st.pyplot(fig_pe, width='stretch')
        plt.close(fig_pe)

    # Row E: Replication + Optimization panel (side-by-side)
    st.divider()
    col_rep, col_opt = st.columns([1, 1], gap="medium")

    with col_rep:
        st.markdown("**Replicated / duplicated storage**")
        rep_entries = replication_report(cfg)
        if not rep_entries:
            st.caption("No duplicated storage detected - every tensor is uniquely sharded.")
        else:
            rep_rows = []
            total_waste = 0
            for e in rep_entries:
                rep_rows.append({
                    "Tensor": e.tensor,
                    "Per-PE (GB)": round(e.per_pe_bytes / 1e9, 3),
                    "Replicated": e.replicated_across,
                    "Copies": e.copies,
                    "Waste/PE (GB)": round(e.wasted_bytes / 1e9, 3),
                })
                total_waste += e.wasted_bytes
            st.dataframe(pd.DataFrame(rep_rows), width='stretch',
                          hide_index=True)
            st.caption(
                f"Total per-PE duplicated bytes: "
                f"**{total_waste/1e9:.2f} GB** "
                f"(vs {mem_layout.budget_bytes/1e9:.1f} GB budget). "
                f"Each 'copy' represents one identical block stored on "
                f"another rank."
            )

    with col_opt:
        st.markdown("**Optimization opportunities**")
        hints = optimization_hints(cfg)
        if not hints:
            st.caption("No obvious improvements at this configuration.")
        else:
            _icons = {"warn": ":warning:", "info": ":bulb:", "good": ":white_check_mark:"}
            _cats = {"space": "SPACE", "comm": "COMM",
                      "compute": "COMPUTE", "layout": "LAYOUT"}
            for h in hints:
                icon = _icons.get(h.severity, ":bulb:")
                cat = _cats.get(h.category, h.category.upper())
                st.markdown(f"{icon} **[{cat}]** {h.message}")

    st.divider()

    # Row F: Per-PE breakdown table (collapsible)
    with st.expander("Per-PE breakdown table", expanded=False):
        rows = []
        for pe_id in range(topo.tp):
            q_heads = _q_heads_for_pe(pe_id, topo.tp, model.h_q)
            kv_heads, kv_note = _kv_heads_for_pe(
                pe_id, topo.tp, model.h_kv, topo.kv_shard_mode,
            )
            bytes_ = _per_pe_bytes(cfg, pe_id)
            weights_bytes = sum(v for k, v in bytes_.items()
                                 if k not in ("KV cache", "Transient"))
            rows.append({
                "PE": pe_id,
                "Q heads": (f"{q_heads[0]}-{q_heads[-1]}"
                            if len(q_heads) > 1 else str(q_heads[0])
                            if q_heads else "-"),
                "KV heads": (str(kv_heads[0]) + (f" ({kv_note})" if kv_note else "")
                             if kv_heads else "-"),
                "W_Q (MB)": round(bytes_["W_Q"] / 1e6, 2),
                "W_K (MB)": round(bytes_["W_K"] / 1e6, 2),
                "W_V (MB)": round(bytes_["W_V"] / 1e6, 2),
                "W_O (MB)": round(bytes_["W_O"] / 1e6, 2),
                "FFN (MB)": round((bytes_["W_gate"] + bytes_["W_up"]
                                    + bytes_["W_down"]) / 1e6, 2),
                "Weights total (GB)": round(weights_bytes / 1e9, 3),
                "KV cache (GB)": round(bytes_["KV cache"] / 1e9, 3),
                "Transient (MB)": round(bytes_["Transient"] / 1e6, 2),
            })
        st.dataframe(pd.DataFrame(rows), width='stretch', hide_index=True)


# ── TAB 2: Memory breakdown ─────────────────────────────────────
with tab_memory:
    mem = compute_memory(cfg)
    layers_per_stage = (model.layers + pp - 1) // pp

    parallelism_str = (
        f"CP={cp}, TP={tp}, PP={pp}, DP={dp}, EP={ep}"
        f" | layers={model.layers}, layers/stage={layers_per_stage}"
        f" | KV mode: {kv_shard_mode}"
    )
    st.caption(f"**Parallelism:** {parallelism_str}")

    # ─── ATTENTION (per-layer) ─────────────────────────────────
    st.subheader(f"Attention weights (per layer)  —  {parallelism_str}")
    attn_rows = attention_weight_rows(cfg)
    display_attn = [{k: v for k, v in r.items() if not k.startswith("_")}
                    for r in attn_rows]
    st.dataframe(pd.DataFrame(display_attn), width='stretch', hide_index=True)

    attn_bytes_1L_global = sum(r["_p_global"] for r in attn_rows) * model.bytes_per_elem
    attn_bytes_1L_pe = sum(r["_p_per_pe"] for r in attn_rows) * model.bytes_per_elem
    attn_bytes_all_global = attn_bytes_1L_global * model.layers
    attn_bytes_all_pe = attn_bytes_1L_pe * layers_per_stage

    ac1, ac2, ac3, ac4 = st.columns(4)
    ac1.metric("Attn / layer (global)", f"{attn_bytes_1L_global/1e6:.2f} MB")
    ac2.metric(f"Attn all {model.layers} layers", f"{attn_bytes_all_global/1e9:.3f} GB")
    ac3.metric("Attn / layer (per PE)", f"{attn_bytes_1L_pe/1e6:.2f} MB")
    ac4.metric(f"Attn per PE ({layers_per_stage} layers/stage)",
                f"{attn_bytes_all_pe/1e9:.3f} GB")

    # ─── FFN / MoE (per-layer) ─────────────────────────────────
    is_moe = "MoE" in (preset.family + preset.note)
    ffn_label = "FFN (activated experts approx)" if is_moe else "FFN"
    st.subheader(f"{ffn_label} weights (per layer)  —  {parallelism_str}")
    ffn_rows = ffn_weight_rows(cfg)
    display_ffn = [{k: v for k, v in r.items() if not k.startswith("_")}
                   for r in ffn_rows]
    st.dataframe(pd.DataFrame(display_ffn), width='stretch', hide_index=True)

    ffn_bytes_1L_global = sum(r["_p_global"] for r in ffn_rows) * model.bytes_per_elem
    ffn_bytes_1L_pe = sum(r["_p_per_pe"] for r in ffn_rows) * model.bytes_per_elem
    ffn_bytes_all_global = ffn_bytes_1L_global * model.layers
    ffn_bytes_all_pe = ffn_bytes_1L_pe * layers_per_stage

    fc1, fc2, fc3, fc4 = st.columns(4)
    fc1.metric("FFN / layer (global)", f"{ffn_bytes_1L_global/1e6:.2f} MB")
    fc2.metric(f"FFN all {model.layers} layers", f"{ffn_bytes_all_global/1e9:.3f} GB")
    fc3.metric("FFN / layer (per PE)", f"{ffn_bytes_1L_pe/1e6:.2f} MB")
    fc4.metric(f"FFN per PE ({layers_per_stage} layers/stage)",
                f"{ffn_bytes_all_pe/1e9:.3f} GB")

    # ─── KV cache (per-layer) ─────────────────────────────────
    st.subheader(f"KV cache (per layer, S_kv = {s_kv:,})  —  {parallelism_str}")
    kv_rows = kv_cache_rows(cfg)
    display_kv = [{k: v for k, v in r.items() if not k.startswith("_")}
                  for r in kv_rows]
    st.dataframe(pd.DataFrame(display_kv), width='stretch', hide_index=True)

    kv_bytes_1L_global = sum(r["_p_global"] for r in kv_rows) * model.bytes_per_elem
    kv_bytes_1L_pe = sum(r["_p_per_pe"] for r in kv_rows) * model.bytes_per_elem
    kv_bytes_all_global = kv_bytes_1L_global * model.layers
    kv_bytes_all_pe = kv_bytes_1L_pe * layers_per_stage

    kc1, kc2, kc3, kc4 = st.columns(4)
    kc1.metric("KV / layer (global)", f"{kv_bytes_1L_global/1e6:.2f} MB")
    kc2.metric(f"KV all {model.layers} layers", f"{kv_bytes_all_global/1e9:.3f} GB")
    kc3.metric("KV / layer (per PE)", f"{kv_bytes_1L_pe/1e6:.2f} MB")
    kc4.metric(f"KV per PE ({layers_per_stage} layers/stage)",
                f"{kv_bytes_all_pe/1e9:.3f} GB")

    st.divider()

    # ─── FULL-MODEL TOTALS ───────────────────────────────────
    st.subheader("Full-model totals (all layers)")
    total_w_global = attn_bytes_all_global + ffn_bytes_all_global
    total_all_global = total_w_global + kv_bytes_all_global
    total_per_pe = attn_bytes_all_pe + ffn_bytes_all_pe + kv_bytes_all_pe
    tc1, tc2, tc3, tc4 = st.columns(4)
    tc1.metric("Weights (all layers, global)", f"{total_w_global/1e9:.2f} GB")
    tc2.metric("KV (all layers, global)", f"{kv_bytes_all_global/1e9:.2f} GB")
    tc3.metric("Weights + KV (global)", f"{total_all_global/1e9:.2f} GB")
    tc4.metric("Per-PE total (W+KV, layers/stage)",
                f"{total_per_pe/1e9:.2f} GB")

    # ─── Per-PE memory pie ───────────────────────────────────
    st.subheader("Per-PE memory footprint (with transient + slack)")
    pie_col, info_col = st.columns([1, 1])
    with pie_col:
        labels = ["Weights", "KV cache", "Transient", "Slack"]
        vals = [mem.weights_bytes, mem.kv_cache_bytes,
                mem.transient_bytes, mem.slack_bytes]
        colors = ["#3a86ff", "#8338ec", "#ffbe0b", "#adb5bd"]

        fig_mem, ax_mem = plt.subplots(figsize=(6, 6))
        ax_mem.pie(vals,
                   labels=[f"{l}\n{v/1e9:.2f} GB" for l, v in zip(labels, vals)],
                   colors=colors, autopct="%1.1f%%", startangle=90)
        ax_mem.set_title(f"Per-PE (budget: {mem.budget_bytes/1e9:.1f} GB)")
        st.pyplot(fig_mem, width='stretch')
        plt.close(fig_mem)
    with info_col:
        if mem.over_budget:
            st.error(f"OVER BUDGET by "
                     f"{(mem.used_bytes - mem.budget_bytes)/1e9:.2f} GB")
        else:
            st.success(f"{mem.slack_bytes/1e9:.2f} GB slack")

        st.markdown(
            f"- Weights: **{mem.weights_bytes/1e9:.2f} GB** "
            f"(layers/PP = {layers_per_stage})\n"
            f"- KV cache: **{mem.kv_cache_bytes/1e9:.2f} GB** "
            f"(S_local = {topo.s_local:,})\n"
            f"- Transient: {mem.transient_bytes/1e9:.2f} GB\n"
            f"- Slack: {mem.slack_bytes/1e9:.2f} GB"
        )


# ── TAB 3: Per-stage latency ────────────────────────────────────
with tab_stages:
    st.subheader("Per-stage attention latency")
    stages = all_stages(cfg)

    rows = []
    for s in stages:
        rows.append({
            "Stage": s.name,
            "Bound": s.bound,
            "Compute (us)": round(s.compute_s * 1e6, 3),
            "Memory (us)": round(s.memory_s * 1e6, 3),
            "Comm (us)": round(s.comm_s * 1e6, 3),
            "Visible (us)": round(s.visible_s * 1e6, 3),
            "Formula": s.formula,
        })
    df = pd.DataFrame(rows)
    st.dataframe(df, width='stretch', hide_index=True)

    total_visible_s = sum(s.visible_s for s in stages)
    total_layers = total_visible_s * model.layers
    st.markdown(
        f"**Per-layer attention (analytical): {total_visible_s*1e6:.2f} us**  \n"
        f"**Full-model attention ({model.layers} layers): "
        f"{total_layers*1e6:.2f} us = {total_layers*1e3:.2f} ms**"
    )

    st.subheader("Stage cost breakdown")
    fig_bar, ax_bar = plt.subplots(figsize=(12, 5))
    names = [s.name for s in stages]
    cmp_vals = [s.compute_s * 1e6 for s in stages]
    mem_vals = [s.memory_s * 1e6 for s in stages]
    comm_vals = [s.comm_s * 1e6 for s in stages]

    x = range(len(names))
    ax_bar.bar(x, cmp_vals, label="Compute", color="#3a86ff")
    ax_bar.bar(x, mem_vals, bottom=cmp_vals, label="Memory", color="#ffbe0b")
    ax_bar.bar(x, comm_vals,
               bottom=[c + m for c, m in zip(cmp_vals, mem_vals)],
               label="Comm", color="#d90429")
    ax_bar.set_xticks(x)
    ax_bar.set_xticklabels(names, rotation=45, ha="right", fontsize=8)
    ax_bar.set_ylabel("Time (us)")
    ax_bar.set_title(f"Stage breakdown - mode={mode}, T_q={topo.T_q}")
    ax_bar.legend()
    ax_bar.grid(axis="y", alpha=0.3)
    plt.tight_layout()
    st.pyplot(fig_bar, width='stretch')
    plt.close(fig_bar)


# ── TAB 4: Save & compare ───────────────────────────────────────
def _snapshot_current_config():
    """Capture a compact snapshot of the current cfg + key metrics."""
    _stages_attn = all_stages(cfg)
    _stages_ffn = all_ffn_stages(cfg)
    attn_us = sum(s.visible_s for s in _stages_attn) * 1e6
    ffn_us = sum(s.visible_s for s in _stages_ffn) * 1e6
    mem = compute_memory(cfg)
    # Per-stage latencies keyed by short prefix (S1, S2, ..., C1, F1, ..., CF1).
    # Value = seconds (raw), so display can format uniformly across snapshots.
    _by_prefix_attn = {s.name.split()[0]: s.visible_s for s in _stages_attn}
    _by_prefix_ffn = {s.name.split()[0]: s.visible_s for s in _stages_ffn}
    # Preserve the full descriptive name too so a hover/second column can
    # show what the abbreviation means for this particular snapshot.
    _fullname_attn = {s.name.split()[0]: s.name for s in _stages_attn}
    _fullname_ffn = {s.name.split()[0]: s.name for s in _stages_ffn}
    return {
        "model": model.name,
        "mode": topo.mode,
        "s_kv": topo.s_kv,
        "cp": topo.cp, "tp": topo.tp,
        "pp": topo.pp, "dp": topo.dp, "ep": topo.ep,
        "cp_placement": topo.cp_placement,
        "tp_placement": topo.tp_placement,
        "ffn_scope": topo.ffn_shard_scope,
        "kv_mode": topo.kv_shard_mode,
        "sip_topology": topo.sip_topology,
        "cp_ring_variant": topo.cp_ring_variant,
        "pes": topo.total_pes,
        "cubes": topo.cubes_used,
        "sips": topo.sips_used,
        "pe_hbm_gb": machine.pe_hbm_gb,
        "peak_tflops": machine.peak_tflops_f16,
        "bw_hbm_gbs": machine.bw_hbm_gbs,
        "bw_inter_gbs": machine.bw_inter_gbs,
        "bw_intersip_gbs": machine.bw_intersip_gbs,
        "weights_gb":   round(mem.weights_bytes / 1e9, 3),
        "kv_gb":        round(mem.kv_cache_bytes / 1e9, 3),
        "transient_gb": round(mem.transient_bytes / 1e9, 3),
        "used_gb":      round(mem.used_bytes / 1e9, 3),
        "slack_gb":     round(mem.slack_bytes / 1e9, 3),
        "over_budget":  mem.over_budget,
        "attn_us":      round(attn_us, 3),
        "ffn_us":       round(ffn_us, 3),
        "layer_us":     round(attn_us + ffn_us, 3),
        "stages_attn":   _by_prefix_attn,
        "stages_ffn":    _by_prefix_ffn,
        "stage_names_attn": _fullname_attn,
        "stage_names_ffn":  _fullname_ffn,
    }


with tab_compare:
    st.markdown(
        "Save the current configuration under a name, then load it back "
        "later or compare multiple saved configs side-by-side."
    )
    if "saved_cfgs" not in st.session_state:
        st.session_state["saved_cfgs"] = {}  # name -> snapshot dict

    # Default name: next free config1, config2, ...
    def _next_config_name():
        i = 1
        while f"config{i}" in st.session_state.get("saved_cfgs", {}):
            i += 1
        return f"config{i}"

    sv_c1, sv_c2, sv_c3 = st.columns([3, 1, 1])
    with sv_c1:
        _save_name = st.text_input(
            "Config name",
            value=_next_config_name(),
            key=f"cfg_save_name_{len(st.session_state.get('saved_cfgs', {}))}",
        )
    with sv_c2:
        if st.button("Save current", width='stretch'):
            if _save_name.strip():
                st.session_state["saved_cfgs"][_save_name.strip()] = \
                    _snapshot_current_config()
                st.success(f"Saved '{_save_name.strip()}'")
                st.rerun()
            else:
                st.warning("Give it a name first.")
    with sv_c3:
        if st.button("Clear all", width='stretch'):
            st.session_state["saved_cfgs"] = {}
            st.rerun()

    _saved_all = st.session_state.get("saved_cfgs", {})
    if not _saved_all:
        st.info("No saved configurations yet.")
    else:
        st.markdown(f"**{len(_saved_all)} saved configuration(s)** - "
                     f"pick which to compare below")

        # Pick which configs to include in the comparison (default all)
        _picked = st.multiselect(
            "Compare",
            options=list(_saved_all.keys()),
            default=list(_saved_all.keys()),
            key="compare_picked",
        )
        _saved = {n: _saved_all[n] for n in _picked}

        # Delete individual
        with st.expander("Manage (delete individual)", expanded=False):
            for _n in list(_saved_all.keys()):
                d1, d2, d3 = st.columns([5, 1, 1])
                d1.text(_n)
                if d2.button("Rename", key=f"rn_btn_{_n}"):
                    st.session_state[f"rn_active_{_n}"] = True
                if d3.button("Delete", key=f"del_{_n}"):
                    del st.session_state["saved_cfgs"][_n]
                    st.rerun()
                if st.session_state.get(f"rn_active_{_n}"):
                    new_name = st.text_input(
                        f"New name for '{_n}'", value=_n,
                        key=f"rn_txt_{_n}",
                    )
                    if st.button("Confirm rename", key=f"rn_ok_{_n}"):
                        if new_name and new_name != _n and new_name not in _saved_all:
                            st.session_state["saved_cfgs"][new_name] = \
                                st.session_state["saved_cfgs"].pop(_n)
                            st.session_state[f"rn_active_{_n}"] = False
                            st.rerun()

        if not _saved:
            st.info("No configurations selected. Tick some above to compare.")
            st.stop()

        # Comparison table: rows = metrics, columns = config names
        _key_order = [
            ("Model / workload", ["model", "mode", "s_kv"]),
            ("Parallelism", ["cp", "tp", "pp", "dp", "ep",
                              "cp_placement", "tp_placement",
                              "ffn_scope", "kv_mode", "cp_ring_variant",
                              "sip_topology"]),
            ("Physical", ["pes", "cubes", "sips"]),
            ("Hardware knobs", ["pe_hbm_gb", "peak_tflops",
                                 "bw_hbm_gbs", "bw_inter_gbs",
                                 "bw_intersip_gbs"]),
            ("Memory / PE (GB)", ["weights_gb", "kv_gb", "transient_gb",
                                    "used_gb", "slack_gb", "over_budget"]),
            ("Per-layer latency (us)", ["attn_us", "ffn_us", "layer_us"]),
        ]

        # Build one wide DataFrame - metrics x saved names
        _names = list(_saved.keys())
        _rows_out = []
        for _sec, _keys in _key_order:
            _rows_out.append({"metric": f"= {_sec}",
                              **{n: "" for n in _names}})
            for k in _keys:
                _rows_out.append({
                    "metric": k,
                    **{n: _saved[n].get(k, "-") for n in _names},
                })
        _df_cmp = pd.DataFrame(_rows_out)

        # Highlight best (lowest) for latency + used_gb rows, worst (over_budget True) red.
        def _cmp_style(row):
            styles = [""] * len(row)
            m = row["metric"]
            vals = row[1:]
            try:
                nums = {k: float(v) for k, v in vals.items()
                        if isinstance(v, (int, float)) and not isinstance(v, bool)}
            except Exception:
                nums = {}
            if m in ("layer_us", "attn_us", "ffn_us", "used_gb",
                      "weights_gb", "kv_gb"):
                if nums:
                    best = min(nums, key=nums.get)
                    for i, (k, v) in enumerate(vals.items(), start=1):
                        if k == best:
                            styles[i] = "background-color: #d4edda; font-weight:600;"
            if m == "slack_gb" and nums:
                best = max(nums, key=nums.get)
                for i, (k, v) in enumerate(vals.items(), start=1):
                    if k == best:
                        styles[i] = "background-color: #d4edda; font-weight:600;"
            if m == "over_budget":
                for i, (k, v) in enumerate(vals.items(), start=1):
                    if v is True or v == "True":
                        styles[i] = "background-color: #f8d7da; font-weight:600;"
            if m.startswith("= "):
                styles = ["background-color: #e9ecef; font-weight:600;"] * len(row)
            return styles

        _styled_cmp = _df_cmp.style.apply(_cmp_style, axis=1)
        st.dataframe(_styled_cmp, width='stretch', hide_index=True,
                      row_height=28)
        st.caption(
            "Green cell = best value in that row (lowest latency / memory, "
            "highest slack). Red = over-budget. Grey rows = section headers."
        )

        # ── Per-stage side-by-side comparison ────────────────────
        _fmt = _fmt_time  # reuse the ns/us/ms auto-formatter
        _short_label = {
            "S1": "S1 RMSNorm", "S2": "S2 W_Q", "S3": "S3 W_K+W_V",
            "S4": "S4 KV append", "S5": "S5 Q.K^T", "S6": "S6 softmax",
            "S7": "S7 P.V", "S8": "S8 merge (+ AR if decode)",
            "S9": "S9 norm O/l", "S10": "S10 W_O",
            "C1": "C1 CP ring (prefill only)",
            "C2": "C2 TP AR (W_O)",
            "C3": "C3 Score AR (head-split)",
            "F1": "F1 RMSNorm", "F2": "F2 W_gate", "F3": "F3 W_up",
            "F4": "F4 SwiGLU", "F5": "F5 W_down", "CF1": "CF1 FFN AR",
        }

        def _stage_rows_across(prefix_order, key):
            """Build one row per stage prefix, with each saved config as a col."""
            rows = []
            for pfx in prefix_order:
                row = {"stage": _short_label.get(pfx, pfx)}
                sec_values = {}
                for n, snap in _saved.items():
                    v = snap.get(key, {}).get(pfx)
                    if v is None:
                        row[n] = "-"
                    else:
                        row[n] = _fmt(v)
                        sec_values[n] = v
                rows.append(row)
            return rows

        def _stage_style_row(row):
            styles = [""] * len(row)
            # Try to find min seconds — reconstruct from formatted strings
            # is fragile; recompute from saved secs directly.
            pfx = row["stage"].split()[0]
            secs = {}
            for n, snap in _saved.items():
                key = "stages_attn" if pfx in _attn_prefixes else "stages_ffn"
                v = snap.get(key, {}).get(pfx)
                if v is not None and v > 0:
                    secs[n] = v
            if secs:
                best = min(secs, key=secs.get)
                for i, name in enumerate(row.index[1:], start=1):
                    if name == best:
                        styles[i] = "background-color: #d4edda; font-weight:600;"
            return styles

        _attn_prefixes = ["S1", "S2", "S3", "S4", "S5", "S6", "S7",
                          "S8", "S9", "S10", "C1", "C2", "C3"]
        _ffn_prefixes = ["F1", "F2", "F3", "F4", "F5", "CF1"]

        st.markdown("**Attention per-stage latency (side-by-side)**")
        _df_attn = pd.DataFrame(_stage_rows_across(_attn_prefixes, "stages_attn"))
        st.dataframe(_df_attn.style.apply(_stage_style_row, axis=1),
                      width='stretch', hide_index=True, row_height=28)

        st.markdown("**FFN per-stage latency (side-by-side)**")
        _df_ffn = pd.DataFrame(_stage_rows_across(_ffn_prefixes, "stages_ffn"))
        st.dataframe(_df_ffn.style.apply(_stage_style_row, axis=1),
                      width='stretch', hide_index=True, row_height=28)

        st.caption(
            "Times auto-scale (ns/us/ms). Green cell = fastest for that stage. "
            "'-' means the stage is absent for that config (e.g. C1 in "
            "decode, since the O/m/l all-reduce is folded into S8)."
        )


# ── TAB (Auto Suggest Parallelism) — one tab, two sweep buttons ──


def _render_auto_explore_tab():
    """Render the Auto Suggest Parallelism tab body.

    Two sweep buttons — Attention-only and Attn+FFN/MoE. Each caches its
    result under its own session_state key. Whichever button was last
    clicked drives the display below; both caches persist so re-clicking
    either without changing model/workload just re-shows the cached view.
    """
    from tests.analytical_visualization.auto_explore import (
        compute_parallelism_sensitivity,
        run_auto_explore,
    )

    st.markdown(
        "Sweep every valid parallelism config (CP, TP, PP, DP, kv_shard, "
        "ffn_scope, tp_placement, cp_placement, cp_ring) for the currently-"
        "selected model + workload and extract the 3D Pareto frontier on "
        "(latency ↓, PEs ↓, efficiency ↑). Throughput = 1/latency for a "
        "single request so it collapses with latency and is shown for info only."
    )
    st.caption(
        "**Attention only**: drop FFN / MoE from the summed latency — "
        "useful when isolating attention-kernel tuning (matches what our "
        "kernbench attention sim measures).  \n"
        "**Attn + FFN/MoE**: sum every stage — the full per-token cost."
    )

    st.markdown(
        f"**Context:** {model.name}  |  S_kv=**{s_kv:,}**  |  mode=**{mode}**  |  "
        f"per-PE HBM = **{machine.pe_hbm_gb:.1f} GB**"
    )

    _b1, _b2, _b3, _b4 = st.columns([1, 1, 1, 1])
    with _b1:
        _run_attn = st.button("Run sweep — Attention",
                               type="primary", width='stretch',
                               key="_auto_run_attn")
    with _b2:
        _run_ffn = st.button("Run sweep — FFN/MoE",
                              type="primary", width='stretch',
                              key="_auto_run_ffn")
    with _b3:
        _run_full = st.button("Run sweep — Attn + FFN/MoE",
                               type="primary", width='stretch',
                               key="_auto_run_full")

    _scope_meta = {
        "attn": (True,  False, "Attention only"),
        "ffn":  (False, True,  "FFN / MoE only"),
        "full": (True,  True,  "Attn + FFN/MoE"),
    }
    _button_click = {"attn": _run_attn, "ffn": _run_ffn, "full": _run_full}

    for _scope, _clicked in _button_click.items():
        if not _clicked:
            continue
        _ia, _ff, _lbl = _scope_meta[_scope]
        with st.spinner(f"Sweeping ~28k configs ({_lbl})..."):
            _r = run_auto_explore(
                model, machine, s_kv=s_kv, mode=mode,
                include_attention=_ia, include_ffn=_ff,
            )
        st.session_state[f"_auto_explore_result_{_scope}"] = _r
        st.session_state[f"_auto_explore_ctx_{_scope}"] = (
            model.name, s_kv, mode, machine.pe_hbm_gb,
        )
        st.session_state["_auto_explore_active"] = _scope

    _active = st.session_state.get("_auto_explore_active")
    if _active is None:
        st.info(
            "Click one of the run buttons above to enumerate valid "
            "parallelism configurations and rank them on the Pareto "
            "frontier. Takes ~5–10 s per scope. All button results are "
            "cached so you can flip between scopes without re-running."
        )
        return

    # Downstream code expects _suffix, _label, include_attention, include_ffn,
    # _res, _ctx_key.
    include_attention, include_ffn, _label = _scope_meta[_active]
    _suffix = f"_{_active}"
    _result_key = f"_auto_explore_result{_suffix}"
    _ctx_key = f"_auto_explore_ctx{_suffix}"
    _res = st.session_state.get(_result_key)
    if _res is None:
        st.info(f"Click **Run sweep — {_label.split(' ')[0]}** to compute this view.")
        return

    st.markdown(f"### Currently viewing: **{_label}**")

    # Warn if the cached result is stale relative to the current config.
    _cached_ctx = st.session_state.get(_ctx_key)
    _cur_ctx = (model.name, s_kv, mode, machine.pe_hbm_gb)
    if _cached_ctx != _cur_ctx:
        st.warning(
            f"Cached {_label} result is from a different model/workload. "
            f"Click **Run sweep — {_label.split(' ')[0]}** to refresh."
        )

    _m1, _m2, _m3, _m4 = st.columns(4)
    _m1.metric("Enumerated", f"{_res.total_enumerated:,}")
    _m2.metric("Feasible", f"{_res.total_feasible:,}")
    _m3.metric("Pareto configs", len(_res.pareto_scores))
    if _res.pareto_scores:
        _best = min(_res.pareto_scores, key=lambda s: s.total_latency_ns)
        _m4.metric("Best latency", f"{_best.latency_ms:.2f} ms")

    if not _res.pareto_scores:
        st.error(
            "No feasible configs — every parallelism choice exceeds "
            "per-PE HBM. Try raising `pe_hbm_gb` in the sidebar or "
            "reducing `s_kv`."
        )
        return

    # ── Pareto scatter: latency vs PEs, coloured by efficiency ─
    import matplotlib.pyplot as _plt
    _fig, _axs = _plt.subplots(1, 2, figsize=(10, 3.5))

    _feas = [s for s in _res.all_scores
             if s.fits_memory and s.placement_valid]
    _axs[0].scatter(
        [s.pes_used for s in _feas],
        [s.latency_ms for s in _feas],
        c="lightgrey", s=8, alpha=0.4, label="feasible",
    )
    _pareto_sorted = sorted(_res.pareto_scores,
                             key=lambda s: s.pes_used)
    _sc = _axs[0].scatter(
        [s.pes_used for s in _pareto_sorted],
        [s.latency_ms for s in _pareto_sorted],
        c=[s.efficiency_score for s in _pareto_sorted],
        cmap="viridis", s=60, edgecolors="black", linewidths=0.8,
        label="Pareto",
    )
    _axs[0].plot(
        [s.pes_used for s in _pareto_sorted],
        [s.latency_ms for s in _pareto_sorted],
        "k--", linewidth=0.8, alpha=0.5,
    )
    _axs[0].set_xlabel("PEs used"); _axs[0].set_ylabel("Latency (ms)")
    _axs[0].set_xscale("log", base=2)
    _axs[0].set_yscale("log")
    _axs[0].set_title("Pareto: latency vs PEs")
    _axs[0].grid(alpha=0.3)
    _axs[0].legend(loc="upper right", fontsize=8)
    _fig.colorbar(_sc, ax=_axs[0], label="efficiency")

    _pareto_by_lat = sorted(_res.pareto_scores,
                             key=lambda s: s.total_latency_ns)
    _axs[1].scatter(
        [s.hbm_utilization * 100 for s in _pareto_by_lat],
        [s.latency_ms for s in _pareto_by_lat],
        c=[s.pes_used for s in _pareto_by_lat],
        cmap="plasma", s=60, edgecolors="black", linewidths=0.8,
    )
    _axs[1].set_xlabel("HBM utilization (%)")
    _axs[1].set_ylabel("Latency (ms)")
    _axs[1].set_yscale("log")
    _axs[1].set_title("HBM usage vs latency (Pareto)")
    _axs[1].grid(alpha=0.3)
    _sc2 = _axs[1].collections[0]
    _fig.colorbar(_sc2, ax=_axs[1], label="PEs")

    _fig.tight_layout()
    st.pyplot(_fig, width='stretch')
    _plt.close(_fig)

    # ── Pareto table ────────────────────────────────────────────
    st.markdown("**Pareto configurations** (sorted by latency)")
    _rows = []
    for i, s in enumerate(_pareto_by_lat):
        _rows.append({
            "#": i,
            "CP": s.cp, "TP": s.tp, "PP": s.pp, "DP": s.dp,
            "kv": s.kv_shard_mode,
            "ffn": s.ffn_shard_scope,
            "tp_place": s.tp_placement,
            "cp_place": s.cp_placement,
            "cp_ring": s.cp_ring_variant,
            "lat (ms)": round(s.latency_ms, 3),
            "eff": round(s.efficiency_score, 4),
            "PEs": s.pes_used,
            "SIPs": s.sips_used,
            "HBM %": round(s.hbm_utilization * 100, 1),
        })
    _df = pd.DataFrame(_rows)
    st.dataframe(_df, width='stretch', hide_index=True)

    # ── Parallelism sensitivity subplots ──────────────────────
    st.markdown(
        "**Parallelism sensitivity** — for the *baseline* config picked "
        "below, vary each knob one at a time (holding others fixed). "
        "Solid line = fits memory; red X = infeasible."
    )
    _sens_c1, _sens_c2 = st.columns([1, 4])
    with _sens_c1:
        _sens_row = st.number_input(
            "Baseline row #",
            min_value=0, max_value=len(_pareto_by_lat) - 1,
            value=0, step=1, key=f"_auto_sens_row{_suffix}",
            help="Which Pareto row is the sensitivity computed around?",
        )
    _baseline_for_sens = _pareto_by_lat[int(_sens_row)]
    _sens_rows = compute_parallelism_sensitivity(
        _baseline_for_sens, model, machine, s_kv=s_kv, mode=mode,
        include_attention=include_attention, include_ffn=include_ffn,
    )

    _fig_sens, _axes_sens = _plt.subplots(1, 5, figsize=(14, 3.2))
    for _ax, _row in zip(_axes_sens, _sens_rows):
        _fit_vals, _fit_lats = [], []
        _bad_vals = []
        for _v, _lat_ns, _fit in zip(_row.values, _row.latencies_ns,
                                      _row.fits_flags):
            if _lat_ns != _lat_ns:
                continue
            if _fit:
                _fit_vals.append(_v)
                _fit_lats.append(_lat_ns / 1e6)
            else:
                _bad_vals.append((_v, _lat_ns / 1e6))
        if _fit_vals:
            _ax.plot(_fit_vals, _fit_lats, "o-", color="#2c3e50",
                      linewidth=1.4, markersize=4)
        if _bad_vals:
            _ax.scatter(
                [v for v, _ in _bad_vals],
                [l for _, l in _bad_vals],
                marker="x", color="#c0392b", s=40, alpha=0.6,
                label="no fit",
            )
        _ax.axvline(_row.baseline_value, color="#3498db",
                     linestyle=":", linewidth=1.0)
        _ax.set_xscale("log")
        _ax.set_yscale("log")
        _ax.set_title(
            f"{_row.knob.upper()}  (baseline={_row.baseline_value})",
            fontsize=10,
        )
        _ax.set_xlabel("value")
        _ax.set_ylabel("lat (ms)")
        _ax.grid(alpha=0.3)
        if _bad_vals:
            _ax.legend(loc="best", fontsize=7)
    _fig_sens.tight_layout()
    st.pyplot(_fig_sens, width='stretch')
    _plt.close(_fig_sens)

    # ── Load-into-sidebar ───────────────────────────────────────
    st.markdown("**Load a Pareto config into the main sidebar sliders**")
    _lc1, _lc2 = st.columns([1, 3])
    with _lc1:
        _pick = st.number_input(
            "Row #", min_value=0, max_value=len(_pareto_by_lat) - 1,
            value=0, step=1, key=f"_auto_pick_row{_suffix}",
        )
    with _lc2:
        if st.button("Load into sidebar", type="secondary",
                      key=f"_auto_load_btn{_suffix}"):
            _s = _pareto_by_lat[int(_pick)]
            st.session_state["cp"] = _s.cp
            st.session_state["tp"] = _s.tp
            st.session_state["pp"] = _s.pp
            st.session_state["dp"] = _s.dp
            st.session_state["tp_placement"] = _s.tp_placement
            st.session_state["cp_placement"] = _s.cp_placement
            st.session_state["cp_ring_variant"] = _s.cp_ring_variant
            st.session_state["kv_mode"] = _s.kv_shard_mode
            _tp, _cp, _dp = _s.tp, _s.cp, _s.dp
            _ffn_label_map = {
                "TP":       f"TP only (div={max(1, _tp)})",
                "TP+CP":    f"TP*CP (div={max(1, _tp * _cp)})",
                "TP+CP+DP": f"TP*CP*DP (div={max(1, _tp * _cp * _dp)})",
            }
            st.session_state["ffn_scope_label"] = \
                _ffn_label_map[_s.ffn_shard_scope]
            st.success(
                f"Loaded row {_pick}: CP={_s.cp} TP={_s.tp} PP={_s.pp} "
                f"DP={_s.dp} → latency {_s.latency_ms:.2f} ms, "
                f"{_s.pes_used} PEs. Flip to another tab to see the "
                f"full breakdown."
            )
            st.rerun()


with tab_auto:
    _render_auto_explore_tab()


# ── TAB (Auto Hardware) — one tab, two sweep buttons ─────────────
def _render_auto_hardware_tab():
    """Render the Auto Hardware tab body.

    Two sweep buttons — Attention-only and Attn+FFN/MoE. Each caches its
    result. Whichever button was last clicked drives the display; both
    caches persist for fast toggling."""
    from tests.analytical_visualization.auto_hardware import (
        _HW_KNOB_DEFAULTS, joint_explore,
    )

    st.markdown(
        "Sweep hardware knobs (PE HBM, HBM BW, TFLOPs, PE↔PE / D2D / C2C "
        "interconnect BW) **together** with parallelism. For each hardware "
        "candidate, we pick the latency-minimum parallelism that fits, then "
        "Pareto-rank the joint (latency ↓, hardware cost ↓) space. Also "
        "computes per-knob sensitivity: which HW investment gives the "
        "biggest speedup?"
    )
    st.caption(
        "**Attention only**: FFN / MoE stages are excluded from the summed "
        "latency (both in the joint search and sensitivity).  \n"
        "**Attn + FFN/MoE**: full per-token cost."
    )

    hx_l, hx_m = st.columns([1, 1])
    with hx_l:
        st.markdown(
            f"**Context:** {model.name}  |  S_kv=**{s_kv:,}**  |  mode=**{mode}**"
        )
    with hx_m:
        _depth = st.radio(
            "Sweep depth",
            options=["two_stage", "balanced", "coarse"],
            index=1, key="_hw_depth", horizontal=True,
            help=("two_stage: 1 HW candidate (defaults) + autosuggest "
                  "parallelism, ~1s. Great for sensitivity alone.\n"
                  "balanced (default): 64 HW × ~2k parallelism, ~10-20s.\n"
                  "coarse: 729 HW × ~2k parallelism, ~2-5 min."),
        )

    _bh1, _bh2, _bh3, _bh4 = st.columns([1, 1, 1, 1])
    with _bh1:
        _run_hw_attn = st.button("Run joint sweep — Attention",
                                   type="primary", width='stretch',
                                   key="_hw_run_attn")
    with _bh2:
        _run_hw_ffn = st.button("Run joint sweep — FFN/MoE",
                                  type="primary", width='stretch',
                                  key="_hw_run_ffn")
    with _bh3:
        _run_hw_full = st.button("Run joint sweep — Attn + FFN/MoE",
                                   type="primary", width='stretch',
                                   key="_hw_run_full")

    _scope_meta_hw = {
        "attn": (True,  False, "Attention only"),
        "ffn":  (False, True,  "FFN / MoE only"),
        "full": (True,  True,  "Attn + FFN/MoE"),
    }
    _button_click_hw = {
        "attn": _run_hw_attn, "ffn": _run_hw_ffn, "full": _run_hw_full,
    }

    for _scope, _clicked in _button_click_hw.items():
        if not _clicked:
            continue
        _ia, _ff, _lbl = _scope_meta_hw[_scope]
        with st.spinner(f"Sweeping HW × parallelism ({_depth}, {_lbl})..."):
            _r = joint_explore(model, s_kv=s_kv, mode=mode, depth=_depth,
                                include_attention=_ia, include_ffn=_ff)
        st.session_state[f"_hw_result_{_scope}"] = _r
        st.session_state[f"_hw_ctx_{_scope}"] = (
            model.name, s_kv, mode, _depth,
        )
        st.session_state["_hw_active"] = _scope

    _active = st.session_state.get("_hw_active")
    if _active is None:
        st.info(
            "Click one of the run buttons to explore hardware × parallelism "
            "co-design for this model + workload. All button results are "
            "cached so you can flip between scopes without re-running."
        )
        return

    include_attention, include_ffn, _label = _scope_meta_hw[_active]
    _suffix = f"_{_active}"
    _hw_result_key = f"_hw_result{_suffix}"
    _hw_ctx_key = f"_hw_ctx{_suffix}"
    _hw_res = st.session_state.get(_hw_result_key)
    if _hw_res is None:
        st.info(f"Click **Run joint sweep — {_label.split(' ')[0]}** to compute this view.")
        return

    st.markdown(f"### Currently viewing: **{_label}**")

    _cached_ctx = st.session_state.get(_hw_ctx_key)
    _cur_ctx = (model.name, s_kv, mode, _depth)
    if _cached_ctx != _cur_ctx:
        st.warning(
            f"Cached {_label} result is from a different model/workload/depth. "
            f"Click **Run joint sweep — {_label.split(' ')[0]}** to refresh."
        )

    _hm1, _hm2, _hm3, _hm4 = st.columns(4)
    _hm1.metric("HW candidates", _hw_res.total_hw)
    _hm2.metric("Feasible joint", _hw_res.total_joint)
    _hm3.metric("Pareto configs", len(_hw_res.pareto_scores))
    if _hw_res.pareto_scores:
        _best = min(_hw_res.pareto_scores,
                     key=lambda s: s.total_latency_ns)
        _hm4.metric("Best latency", f"{_best.latency_ms:.2f} ms")

    if not _hw_res.pareto_scores:
        st.error(
            "No feasible joint configurations — every HW+parallelism "
            "combo exceeds per-PE HBM at this workload."
        )
        return
    if True:
            # ── Panel 1: latency vs cost Pareto ────────────────────────
            import matplotlib.pyplot as _plt
            _fig1, _ax1 = _plt.subplots(figsize=(8, 4))
            _feas = _hw_res.all_scores
            _ax1.scatter(
                [s.cost_score for s in _feas],
                [s.latency_ms for s in _feas],
                c="lightgrey", s=15, alpha=0.5, label="feasible",
            )
            _pareto_by_cost = sorted(_hw_res.pareto_scores,
                                      key=lambda s: s.cost_score)
            _sc = _ax1.scatter(
                [s.cost_score for s in _pareto_by_cost],
                [s.latency_ms for s in _pareto_by_cost],
                c=[s.parallelism.pes_used for s in _pareto_by_cost],
                cmap="viridis", s=90, edgecolors="black", linewidths=1.0,
                label="Pareto",
            )
            _ax1.plot(
                [s.cost_score for s in _pareto_by_cost],
                [s.latency_ms for s in _pareto_by_cost],
                "k--", linewidth=1.0, alpha=0.6,
            )
            _ax1.set_xlabel("Hardware cost proxy (unitless — 6.0 = all defaults)")
            _ax1.set_ylabel("Latency (ms)")
            _ax1.set_yscale("log")
            _ax1.set_title("Optimal HW+parallelism: latency vs cost")
            _ax1.grid(alpha=0.3)
            _ax1.legend(loc="upper right", fontsize=9)
            _fig1.colorbar(_sc, ax=_ax1, label="PEs")
            _fig1.tight_layout()
            st.pyplot(_fig1, width='stretch')
            _plt.close(_fig1)

            # ── Panel 2: per-knob sensitivity ──────────────────────────
            st.markdown(
                "**Per-knob sensitivity** (doubling each HW knob from the "
                "latency-optimal baseline). Higher bar = more valuable "
                "investment for this workload."
            )
            _fig2, _ax2 = _plt.subplots(figsize=(8, 3.5))
            _knobs = [r.knob for r in _hw_res.sensitivity]
            _speedups = [r.rel_speedup * 100 for r in _hw_res.sensitivity]
            _colors = ["#2ecc71" if v > 0 else "#95a5a6" for v in _speedups]
            _ax2.barh(_knobs, _speedups, color=_colors,
                       edgecolor="black", linewidth=0.5)
            _ax2.invert_yaxis()  # biggest speedup on top
            _ax2.set_xlabel("Speedup when knob is doubled (%)")
            _ax2.set_title("Sensitivity — where to invest next")
            _ax2.grid(alpha=0.3, axis="x")
            for i, r in enumerate(_hw_res.sensitivity):
                _ax2.text(
                    r.rel_speedup * 100 + 0.5, i,
                    f"{r.baseline_value:g} → {r.doubled_value:g}",
                    va="center", fontsize=8,
                )
            _fig2.tight_layout()
            st.pyplot(_fig2, width='stretch')
            _plt.close(_fig2)

            # ── Panel 3: Pareto table (HW spec + parallelism) ──────────
            st.markdown("**Pareto joint configurations** (sorted by latency)")
            _pareto_by_lat = sorted(_hw_res.pareto_scores,
                                     key=lambda s: s.total_latency_ns)
            _rows = []
            for i, s in enumerate(_pareto_by_lat):
                p = s.parallelism
                h = s.hardware
                _rows.append({
                    "#": i,
                    "lat (ms)": round(s.latency_ms, 3),
                    "cost": round(s.cost_score, 2),
                    "PEs": p.pes_used,
                    "CP": p.cp, "TP": p.tp, "PP": p.pp, "DP": p.dp,
                    "kv": p.kv_shard_mode,
                    "PE HBM": f"{h.pe_hbm_gb:g} GB",
                    "HBM BW": f"{h.bw_hbm_gbs:g} GB/s",
                    "TFLOPs": f"{h.peak_tflops_f16:g}",
                    "PE↔PE": f"{h.bw_intra_gbs:g}",
                    "D2D": f"{h.bw_inter_gbs:g}",
                    "C2C": f"{h.bw_intersip_gbs:g}",
                })
            _df_hw = pd.DataFrame(_rows)
            st.dataframe(_df_hw, width='stretch', hide_index=True)

            # ── Load hardware into sidebar ─────────────────────────────
            st.markdown(
                "**Load a Pareto HW config into the main sidebar** "
                "(overwrites the Hardware→Per-PE and Hardware→Interconnect "
                "sliders; also loads the paired parallelism)."
            )
            _lch1, _lch2 = st.columns([1, 3])
            with _lch1:
                _pick_hw = st.number_input(
                    "Row #", min_value=0,
                    max_value=len(_pareto_by_lat) - 1,
                    value=0, step=1, key=f"_hw_pick_row{_suffix}",
                )
            with _lch2:
                if st.button("Load into sidebar", type="secondary",
                              key=f"_hw_load_btn{_suffix}"):
                    _s = _pareto_by_lat[int(_pick_hw)]
                    _h = _s.hardware
                    _p = _s.parallelism
                    # HW knobs (only load if the value is in the selectbox
                    # options; otherwise leave as-is).
                    _apply_if_in = {
                        "pe_hbm":       ([3.0, 6.0, 12.0, 24.0, 48.0, 96.0],
                                          _h.pe_hbm_gb),
                        "bw_hbm":       ([128.0, 256.0, 512.0, 1024.0, 2048.0],
                                          _h.bw_hbm_gbs),
                        "tflops":       ([2.0, 4.0, 8.0, 16.0, 32.0, 64.0],
                                          _h.peak_tflops_f16),
                        "bw_intra":     ([128.0, 256.0, 512.0, 1024.0, 2048.0,
                                          4096.0], _h.bw_intra_gbs),
                        "bw_inter":     ([32.0, 64.0, 128.0, 256.0, 512.0,
                                          900.0], _h.bw_inter_gbs),
                        "bw_intersip":  ([12.5, 25.0, 50.0, 100.0, 200.0,
                                          400.0], _h.bw_intersip_gbs),
                    }
                    for _key, (_opts, _val) in _apply_if_in.items():
                        if _val in _opts:
                            st.session_state[_key] = _val
                    # Parallelism knobs
                    st.session_state["cp"] = _p.cp
                    st.session_state["tp"] = _p.tp
                    st.session_state["pp"] = _p.pp
                    st.session_state["dp"] = _p.dp
                    st.session_state["tp_placement"] = _p.tp_placement
                    st.session_state["cp_placement"] = _p.cp_placement
                    st.session_state["cp_ring_variant"] = _p.cp_ring_variant
                    st.session_state["kv_mode"] = _p.kv_shard_mode
                    # Ffn_scope_label reconstruction (same as auto_explore tab)
                    _tp2, _cp2, _dp2 = _p.tp, _p.cp, _p.dp
                    _ffn_label_map = {
                        "TP":       f"TP only (div={max(1, _tp2)})",
                        "TP+CP":    f"TP*CP (div={max(1, _tp2 * _cp2)})",
                        "TP+CP+DP": f"TP*CP*DP (div={max(1, _tp2 * _cp2 * _dp2)})",
                    }
                    st.session_state["ffn_scope_label"] = \
                        _ffn_label_map[_p.ffn_shard_scope]
                    st.success(
                        f"Loaded row {_pick_hw}: latency {_s.latency_ms:.2f} "
                        f"ms, cost {_s.cost_score:.2f}, "
                        f"{_p.pes_used} PEs. Flip to another tab to see the "
                        f"full breakdown."
                    )
                    st.rerun()


with tab_hw:
    _render_auto_hardware_tab()
