"""Draw a PE-level view of one CP group (one TP group of cubes).

Shows how weights + KV cache are distributed across the PEs of ONE CP
rank. Attention weights are sharded by head (H_q, H_kv split across TP);
FFN is sharded by dim (ffn_dim / TP / EP); KV cache is sharded by CP
(sequence) × TP (head).

If TP > 8 the group spans multiple cubes — layout draws all of them.
"""
from __future__ import annotations

import matplotlib.pyplot as plt
import matplotlib.patches as patches

from .model_config import FullConfig


PE_COLS = 4
PE_ROWS = 2
PES_PER_CUBE = PE_COLS * PE_ROWS


def layers_per_stage_val(cfg: FullConfig) -> int:
    return (cfg.model.layers + cfg.topo.pp - 1) // cfg.topo.pp


def _q_heads_for_pe(pe_id_in_group: int, tp: int, h_q: int) -> list[int]:
    """Return the list of Q-head indices assigned to this PE."""
    heads_per_pe = h_q / tp
    start = int(round(pe_id_in_group * heads_per_pe))
    end = int(round((pe_id_in_group + 1) * heads_per_pe))
    return list(range(start, end))


def _kv_heads_for_pe(pe_id_in_group: int, tp: int, h_kv: int,
                      kv_shard_mode: str) -> tuple[list[int], str]:
    """KV heads assigned to this PE and a note about replication/split."""
    if h_kv >= tp:
        heads_per_pe = h_kv // tp
        start = pe_id_in_group * heads_per_pe
        end = start + heads_per_pe
        return list(range(start, end)), ""
    # TP > H_kv
    if kv_shard_mode == "replicate":
        rep_factor = tp // h_kv
        head_idx = pe_id_in_group // rep_factor
        return [head_idx], f"replicated x{rep_factor}"
    # split mode
    split_factor = tp // h_kv
    head_idx = pe_id_in_group // split_factor
    part = pe_id_in_group % split_factor
    return [head_idx], f"head-split ({part+1}/{split_factor})"


def _per_pe_bytes(cfg: FullConfig, pe_id_in_group: int) -> dict:
    """Return per-tensor bytes for one PE within a TP group."""
    m = cfg.model
    tp = cfg.topo.tp
    pp = cfg.topo.pp
    ep = max(1, cfg.topo.ep)
    b = m.bytes_per_elem
    layers_per_stage = (m.layers + pp - 1) // pp

    hq_per_pe = m.h_q / tp
    if cfg.topo.kv_shard_mode == "replicate":
        hkv_per_pe = max(1.0, m.h_kv / tp)
    else:
        hkv_per_pe = m.h_kv / tp
    ffn_div = cfg.ffn_shard_divisor * ep
    ffn_per_pe = m.ffn_dim // ffn_div

    per_layer = {
        "W_Q": int(m.hidden * hq_per_pe * m.d_head * b),
        "W_K": int(m.hidden * hkv_per_pe * m.d_head * b),
        "W_V": int(m.hidden * hkv_per_pe * m.d_head * b),
        "W_O": int(hq_per_pe * m.d_head * m.hidden * b),
        "W_gate": int(m.hidden * ffn_per_pe * b),
        "W_up": int(m.hidden * ffn_per_pe * b),
        "W_down": int(ffn_per_pe * m.hidden * b),
    }
    kv_per_layer = int(2 * cfg.topo.s_local * hkv_per_pe * m.d_head * b)

    all_layers = {k: v * layers_per_stage for k, v in per_layer.items()}
    all_layers["KV cache"] = kv_per_layer * layers_per_stage
    # transient estimate
    if cfg.topo.mode == "decode":
        all_layers["Transient"] = int(4 * (m.hidden + hq_per_pe * m.d_head) * b)
    else:
        TILE = 1024
        all_layers["Transient"] = int(2 * hq_per_pe * cfg.topo.T_q * TILE * b
                                       + cfg.topo.T_q * m.hidden * b)
    return all_layers


def draw_pe_layout(cfg: FullConfig, ax=None):
    """Draw one CP group's PEs with per-tensor weight + KV breakdown.

    Each PE shows: PE id, Q heads, KV heads, W_Q / W_K / W_V / W_O in
    MB, FFN (gate+up+down) in MB, KV cache in GB, total in GB.

    TP=8 -> 1 cube. TP=16 -> 2 cubes. TP=32 -> 4 cubes.
    Under placement, n_cubes reflects cubes-per-group (may include CP spill).
    """
    tp = cfg.topo.tp
    # cubes per one cube-level group (= intra-cube dims / PEs-per-cube, rounded up).
    n_cubes = max(1, (cfg.topo.intra_cube_dims + PES_PER_CUBE - 1) // PES_PER_CUBE)

    if ax is None:
        fig, ax = plt.subplots(figsize=(7.5 * n_cubes, 8.5))
    else:
        fig = ax.figure

    cube_w = 7.0
    cube_h = 6.0
    cube_gap = 0.6

    for cube_idx in range(n_cubes):
        x0 = cube_idx * (cube_w + cube_gap)
        y0 = 0

        rect = patches.FancyBboxPatch(
            (x0, y0), cube_w, cube_h,
            boxstyle="round,pad=0.05",
            facecolor="#f8f9fa", edgecolor="#212529", linewidth=1.5,
        )
        ax.add_patch(rect)
        ax.text(x0 + cube_w / 2, y0 + cube_h + 0.1,
                f"Cube {cube_idx}  (PEs {cube_idx*8}-{cube_idx*8+7})",
                ha="center", va="bottom", fontsize=11, fontweight="bold")

        pe_pad_x = 0.15
        pe_pad_y = 0.2
        pe_gap = 0.10
        inner_w = cube_w - 2 * pe_pad_x
        inner_h = cube_h - 2 * pe_pad_y
        pe_w = (inner_w - (PE_COLS - 1) * pe_gap) / PE_COLS
        pe_h = (inner_h - (PE_ROWS - 1) * pe_gap) / PE_ROWS

        # When cp_placement=pe, CP ranks are packed intra-cube along with
        # TP ranks. Color-code by CP rank so it's obvious which PEs belong
        # to which sequence-shard group. Palette wraps beyond 8 CP ranks.
        _cp_place = cfg.topo.cp_placement
        _cp_val = cfg.topo.cp
        _cp_palette = [
            "#f0f7ff",  # baseline blue (also the fallback for cp_placement=cube)
            "#e8f5e9",  # light green
            "#fff3e0",  # light orange
            "#f3e5f5",  # light purple
            "#ffebee",  # light red
            "#e0f7fa",  # light cyan
            "#fce4ec",  # light pink
            "#f9fbe7",  # light lime
        ]
        _cp_edge_palette = [
            "#3a86ff", "#2e7d32", "#ef6c00", "#7b1fa2",
            "#c62828", "#0097a7", "#c2185b", "#9e9d24",
        ]

        for pr in range(PE_ROWS):
            for pc in range(PE_COLS):
                pe_local = pr * PE_COLS + pc
                pe_id_in_group = cube_idx * PES_PER_CUBE + pe_local
                if pe_id_in_group >= cfg.topo.intra_cube_dims:
                    continue
                px = x0 + pe_pad_x + pc * (pe_w + pe_gap)
                py = y0 + pe_pad_y + (PE_ROWS - 1 - pr) * (pe_h + pe_gap)

                # Which (cp_rank, tp_rank) does this PE hold?
                # When cp_placement=pe: cp_rank = pe_id // tp, tp_rank = pe_id % tp.
                # When cp_placement=cube: every PE in this cube is one CP rank
                # (drawn on its own tile), tp_rank = pe_id % tp.
                if _cp_place == "pe" and _cp_val > 1:
                    _cp_rank = pe_id_in_group // tp
                    _tp_rank = pe_id_in_group % tp
                    _fc = _cp_palette[_cp_rank % len(_cp_palette)]
                    _ec = _cp_edge_palette[_cp_rank % len(_cp_edge_palette)]
                    _cp_label = f"  |  CP={_cp_rank}"
                else:
                    _cp_rank = None
                    _tp_rank = pe_id_in_group % tp
                    _fc = "#f0f7ff"
                    _ec = "#3a86ff"
                    _cp_label = ""

                pe_rect = patches.Rectangle(
                    (px, py), pe_w, pe_h,
                    facecolor=_fc, edgecolor=_ec, linewidth=1.2,
                )
                ax.add_patch(pe_rect)

                # Head assignment uses TP rank (not raw pe_id) so that under
                # cp_placement=pe, all CP ranks share the same head split.
                q_heads = _q_heads_for_pe(_tp_rank, tp, cfg.model.h_q)
                kv_heads, kv_note = _kv_heads_for_pe(
                    _tp_rank, tp, cfg.model.h_kv, cfg.topo.kv_shard_mode,
                )
                bytes_ = _per_pe_bytes(cfg, _tp_rank)
                weights_gb = sum(v for k, v in bytes_.items()
                                 if k not in ("KV cache", "Transient")) / 1e9
                kv_gb = bytes_["KV cache"] / 1e9
                trans_mb = bytes_["Transient"] / 1e6
                total_gb = sum(bytes_.values()) / 1e9

                q_str = (f"{q_heads[0]}-{q_heads[-1]}"
                         if len(q_heads) > 1 else str(q_heads[0])
                         if q_heads else "-")
                kv_str = str(kv_heads[0]) if kv_heads else "-"
                if kv_note:
                    kv_str += f"({kv_note})"

                ffn_mb = (bytes_["W_gate"] + bytes_["W_up"]
                          + bytes_["W_down"]) / 1e6
                # Header (PE id and heads) bigger, then per-tensor breakdown
                header = (f"PE {pe_id_in_group}  TP={_tp_rank}{_cp_label}\n"
                          f"Q heads: {q_str}\n"
                          f"KV head: {kv_str}")
                ax.text(px + 0.08, py + pe_h - 0.05, header,
                        ha="left", va="top",
                        fontsize=7.0, fontweight="bold",
                        family="monospace", color="#0d47a1")

                # Per-tensor bytes with explicit division divisors.
                # Divisor is always the PRODUCT of the parallel dims listed
                # (TP*CP means TP times CP, i.e. multiplication).
                tp_d = cfg.topo.tp
                ffn_div = cfg.ffn_shard_divisor
                ffn_scope = cfg.topo.ffn_shard_scope.replace("+", "*")
                kv_div = cfg.topo.cp * cfg.topo.tp
                L = layers_per_stage_val(cfg)
                detail = (
                    f"weights (all {L} layers):\n"
                    f" W_Q /TP={tp_d} : {bytes_['W_Q']/1e6:7.1f}MB\n"
                    f" W_K /TP={tp_d} : {bytes_['W_K']/1e6:7.1f}MB\n"
                    f" W_V /TP={tp_d} : {bytes_['W_V']/1e6:7.1f}MB\n"
                    f" W_O /TP={tp_d} : {bytes_['W_O']/1e6:7.1f}MB\n"
                    f" FFN /({ffn_scope})={ffn_div}: {ffn_mb:7.1f}MB\n"
                    f"runtime:\n"
                    f" KV /(CP*TP)={kv_div}: {kv_gb*1000:7.1f}MB\n"
                    f" trans     : {trans_mb:7.1f}MB\n"
                    f"W total: {weights_gb*1000:7.1f}MB\n"
                    f"TOTAL  : {total_gb*1000:7.1f}MB"
                )
                ax.text(px + pe_w / 2, py + 0.06, detail,
                        ha="center", va="bottom",
                        fontsize=5.8, family="monospace", color="#1a1a1a")

    total_w = n_cubes * cube_w + (n_cubes - 1) * cube_gap
    ax.set_xlim(-0.2, total_w + 0.2)
    ax.set_ylim(-0.3, cube_h + 0.6)
    ax.set_aspect("equal")
    ax.axis("off")

    title = (
        f"Per-PE layout of one CP group  "
        f"(TP={tp}, H_q={cfg.model.h_q}, H_kv={cfg.model.h_kv})"
        f" | KV mode: {cfg.topo.kv_shard_mode}"
    )
    ax.set_title(title, fontsize=11, fontweight="bold")

    return fig
