"""Pictorial view of how each weight/KV tensor is sharded across ranks.

Every tensor is drawn as a rectangle labelled with its (rows, cols)
shape. Grid lines show shard boundaries; one shard (this PE's share)
is highlighted so it's obvious which dim the split is along.

Sharding pattern:
  - W_Q  (hidden, H_q*d_head) : split on **output cols** across TP
  - W_K  (hidden, H_kv*d_head): split on **output cols** across TP
  - W_V  (hidden, H_kv*d_head): split on **output cols** across TP
  - W_O  (H_q*d_head, hidden) : split on **input rows** across TP
  - W_gate/up (hidden, ffn_dim): split on **output cols** across FFN scope
  - W_down (ffn_dim, hidden)   : split on **input rows** across FFN scope
  - K/V cache (S_kv, H_kv*d_head): split on **rows (S axis)** by CP AND
    **cols (head axis)** by TP -> 2D shard grid
"""
from __future__ import annotations

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

from .model_config import FullConfig


_UNSHARDED_FACE = "#e9ecef"
_UNSHARDED_EDGE = "#adb5bd"
_MY_SHARD_FACE = "#3a86ff"
_MY_SHARD_EDGE = "#0057b3"
_TP_LINE_COLOR = "#d90429"     # TP split lines
_CP_LINE_COLOR = "#0f9d58"     # CP split lines
_FFN_LINE_COLOR = "#e37400"    # FFN-scope split lines


def _draw_tensor(ax, x, y, w, h, name, shape_str,
                  row_splits: int = 1, col_splits: int = 1,
                  my_row_shard: int = 0, my_col_shard: int = 0,
                  row_line_color: str = _TP_LINE_COLOR,
                  col_line_color: str = _TP_LINE_COLOR,
                  row_label: str = "", col_label: str = "",
                  note: str = "",
                  cell_annot=None):
    """Draw a tensor with sharding grid.

    cell_annot: optional callable (row, col) -> str returning the physical
    label to render inside shard cell (r, c). If it returns "", cell is not
    annotated.
    """
    """Draw a tensor as a rectangle with grid lines showing shards."""
    # Base rectangle (unsharded background)
    rect = patches.Rectangle(
        (x, y), w, h,
        facecolor=_UNSHARDED_FACE, edgecolor=_UNSHARDED_EDGE, linewidth=1.0,
    )
    ax.add_patch(rect)

    # Highlight this PE's shard
    shard_w = w / max(1, col_splits)
    shard_h = h / max(1, row_splits)
    sx = x + my_col_shard * shard_w
    sy = y + (row_splits - 1 - my_row_shard) * shard_h
    highlight = patches.Rectangle(
        (sx, sy), shard_w, shard_h,
        facecolor=_MY_SHARD_FACE, edgecolor=_MY_SHARD_EDGE, linewidth=1.5,
        alpha=0.8,
    )
    ax.add_patch(highlight)

    # Horizontal grid lines (row splits)
    for r in range(1, row_splits):
        yy = y + r * shard_h
        ax.plot([x, x + w], [yy, yy],
                color=row_line_color, linewidth=1.2, linestyle="--")
    # Vertical grid lines (col splits)
    for c in range(1, col_splits):
        xx = x + c * shard_w
        ax.plot([xx, xx], [y, y + h],
                color=col_line_color, linewidth=1.2, linestyle="--")

    # Per-cell physical annotations (e.g. "cube3 PE5") if provided.
    if cell_annot is not None:
        for r in range(row_splits):
            for c in range(col_splits):
                text = cell_annot(r, c)
                if not text:
                    continue
                cx = x + c * shard_w + shard_w / 2
                cy = y + (row_splits - 1 - r) * shard_h + shard_h / 2
                ax.text(cx, cy, text,
                        ha="center", va="center",
                        fontsize=6.0, color="#0d1b2a",
                        alpha=0.85)

    # Stacked labels ABOVE rectangle: note (italic, top), shape, name (bold, closest)
    if note:
        ax.text(x + w / 2, y + h + 0.55, note,
                ha="center", va="bottom", fontsize=6.5,
                color="#555", style="italic")
    ax.text(x + w / 2, y + h + 0.30, shape_str,
            ha="center", va="bottom", fontsize=7, color="#666")
    ax.text(x + w / 2, y + h + 0.08, name,
            ha="center", va="bottom",
            fontsize=10, fontweight="bold")

    # Split labels BELOW rectangle (stacked if both present)
    if row_label and col_label:
        ax.text(x + w / 2, y - 0.10, row_label,
                ha="center", va="top", fontsize=6.5,
                color=row_line_color, fontweight="bold")
        ax.text(x + w / 2, y - 0.35, col_label,
                ha="center", va="top", fontsize=6.5,
                color=col_line_color, fontweight="bold")
    elif row_label:
        ax.text(x + w / 2, y - 0.10, row_label,
                ha="center", va="top", fontsize=6.5,
                color=row_line_color, fontweight="bold")
    elif col_label:
        ax.text(x + w / 2, y - 0.10, col_label,
                ha="center", va="top", fontsize=6.5,
                color=col_line_color, fontweight="bold")


def _physical_label(cfg: FullConfig, dim: str, rank: int) -> str:
    """Where does rank `rank` of dim (`tp`|`cp`|`ffn`) live physically?
    Returns a short "cubeX PEy" style label based on the placement config.
    """
    topo = cfg.topo
    if dim == "tp":
        if topo.tp_placement == "pe":
            # TP fills PEs within one cube (spilling if tp > 8)
            pe_in_cube = rank % topo.pes_per_cube_hw
            cube = rank // topo.pes_per_cube_hw
            return f"c{cube}p{pe_in_cube}"
        else:  # tp on cube
            return f"cube{rank}"
    if dim == "cp":
        if topo.cp_placement == "pe":
            pe_in_cube = rank % topo.pes_per_cube_hw
            cube = rank // topo.pes_per_cube_hw
            return f"c{cube}p{pe_in_cube}"
        else:  # cp on cube
            return f"cube{rank}"
    return ""


def _kv_cell_annot(cfg: FullConfig):
    """Return an annotation function for KV cache cells: (cp_r, tp_r) -> label."""
    def _ann(r, c):
        cp_lbl = _physical_label(cfg, "cp", r)
        tp_lbl = _physical_label(cfg, "tp", c)
        # Two lines if both are non-trivial; otherwise a single line
        if cp_lbl and tp_lbl:
            return f"{cp_lbl}\n{tp_lbl}"
        return cp_lbl or tp_lbl
    return _ann


def _tp_cell_annot(cfg: FullConfig):
    """Col-parallel tensors (W_Q/W_K/W_V/W_gate/W_up): each col = one TP rank."""
    def _ann(r, c):
        return _physical_label(cfg, "tp", c)
    return _ann


def _tp_row_cell_annot(cfg: FullConfig):
    """Row-parallel tensors split by TP (W_O): each row = one TP rank."""
    def _ann(r, c):
        return _physical_label(cfg, "tp", r)
    return _ann


def draw_tensor_sharding(cfg: FullConfig, ax=None, my_pe: int = 0,
                          show_physical: bool = False):
    """Draw all tensors with their sharding overlays for PE `my_pe`.

    show_physical: annotate each shard cell with the owning cube/PE.
    """
    m = cfg.model
    tp = cfg.topo.tp
    cp = cfg.topo.cp

    # PE's assigned shard index for each dim
    q_shard = my_pe % tp  # W_Q col shard
    kv_shard_col = min(my_pe % tp, m.h_kv - 1) if m.h_kv >= tp else 0
    ffn_shard = my_pe % cfg.ffn_shard_divisor
    kv_row_shard = 0  # per-CP; showing CP=0 view for clarity

    # Enlarged mode gives more room for physical annotations.
    _fig_w, _fig_h = (22, 13) if show_physical else (15, 9)
    _cell_w = 3.4 if show_physical else 2.6
    _cell_h = 3.0 if show_physical else 2.2
    _gap_x = 0.8 if show_physical else 0.55
    _gap_y = 2.0 if show_physical else 1.7

    if ax is None:
        fig, ax = plt.subplots(figsize=(_fig_w, _fig_h))
    else:
        fig = ax.figure

    # Grid layout: 2 rows x 4 cols of tensor rects
    # Attention row: W_Q, W_K, W_V, W_O
    # FFN + KV row: W_gate, W_up, W_down, KV cache
    cell_w = _cell_w
    cell_h = _cell_h
    gap_x = _gap_x
    gap_y = _gap_y

    _kv_ann = _kv_cell_annot(cfg) if show_physical else None
    _tp_col_ann = _tp_cell_annot(cfg) if show_physical else None
    _tp_row_ann = _tp_row_cell_annot(cfg) if show_physical else None

    positions = [
        (0, 1),  # W_Q at row 0
        (1, 1),  # W_K
        (2, 1),  # W_V
        (3, 1),  # W_O
        (0, 0),  # W_gate at row 1
        (1, 0),  # W_up
        (2, 0),  # W_down
        (3, 0),  # KV cache
    ]

    # ── attention tensors ────────────────────────
    tp_txt = f"TP={tp}"
    ffn_txt = f"{cfg.topo.ffn_shard_scope}={cfg.ffn_shard_divisor}"
    cp_txt = f"CP={cp}"

    def px(col): return col * (cell_w + gap_x)
    def py(row): return row * (cell_h + gap_y)

    # W_Q: (hidden, H_q*d_head), split on COL by TP
    _draw_tensor(ax, px(0), py(1), cell_w, cell_h,
                  "W_Q", f"({m.hidden}, {m.h_q * m.d_head})",
                  row_splits=1, col_splits=tp,
                  my_col_shard=q_shard,
                  col_line_color=_TP_LINE_COLOR,
                  col_label=f"cols split by {tp_txt}",
                  note="col-parallel (output dim)",
                  cell_annot=_tp_col_ann)

    # W_K, W_V: (hidden, H_kv*d_head)
    for (name, col_idx) in [("W_K", 1), ("W_V", 2)]:
        splits = min(tp, m.h_kv) if m.h_kv >= tp else tp
        _draw_tensor(ax, px(col_idx), py(1), cell_w, cell_h,
                      name, f"({m.hidden}, {m.h_kv * m.d_head})",
                      row_splits=1, col_splits=splits,
                      my_col_shard=min(q_shard, splits - 1),
                      col_line_color=_TP_LINE_COLOR,
                      col_label=f"cols split by {tp_txt}",
                      note="col-parallel (heads)",
                      cell_annot=_tp_col_ann)

    # W_O: (H_q*d_head, hidden), split on ROW by TP
    _draw_tensor(ax, px(3), py(1), cell_w, cell_h,
                  "W_O", f"({m.h_q * m.d_head}, {m.hidden})",
                  row_splits=tp, col_splits=1,
                  my_row_shard=q_shard,
                  row_line_color=_TP_LINE_COLOR,
                  row_label=f"rows split by {tp_txt}",
                  note="row-parallel (input dim)",
                  cell_annot=_tp_row_ann)

    # ── FFN tensors ──────────────────────────────
    # W_gate, W_up: (hidden, ffn_dim), split on COL by FFN scope
    for (name, col_idx) in [("W_gate", 0), ("W_up", 1)]:
        _draw_tensor(ax, px(col_idx), py(0), cell_w, cell_h,
                      name, f"({m.hidden}, {m.ffn_dim})",
                      row_splits=1, col_splits=cfg.ffn_shard_divisor,
                      my_col_shard=ffn_shard,
                      col_line_color=_FFN_LINE_COLOR,
                      col_label=f"cols split by {ffn_txt}",
                      note="col-parallel")

    # W_down: (ffn_dim, hidden), split on ROW by FFN scope
    _draw_tensor(ax, px(2), py(0), cell_w, cell_h,
                  "W_down", f"({m.ffn_dim}, {m.hidden})",
                  row_splits=cfg.ffn_shard_divisor, col_splits=1,
                  my_row_shard=ffn_shard,
                  row_line_color=_FFN_LINE_COLOR,
                  row_label=f"rows split by {ffn_txt}",
                  note="row-parallel")

    # ── KV cache: (S_kv, H_kv*d_head), 2D shard ─
    kv_row_splits = cp
    kv_col_splits = min(tp, m.h_kv) if m.h_kv >= tp else tp
    _draw_tensor(ax, px(3), py(0), cell_w, cell_h,
                  "KV cache (K and V)",
                  f"({m.h_kv * m.d_head}, S_kv={cfg.topo.s_kv:,})",
                  row_splits=kv_row_splits,
                  col_splits=kv_col_splits if show_physical else 1,
                  my_row_shard=kv_row_shard,
                  row_line_color=_CP_LINE_COLOR,
                  row_label=f"rows: S split by {cp_txt}",
                  col_label=f"cols: heads split by {tp_txt}",
                  col_line_color=_TP_LINE_COLOR,
                  note="2D shard: seq axis (CP) x head axis (TP)",
                  cell_annot=_kv_ann)
    # In non-physical mode, KV was drawn with col_splits=1; overlay TP col
    # splits as dotted lines to hint at the 2D nature. In physical mode
    # the _draw_tensor call already drew proper vertical dashes.
    if not show_physical:
        kv_x = px(3)
        kv_y = py(0)
        shard_col_w = cell_w / max(1, kv_col_splits)
        for c in range(1, kv_col_splits):
            xx = kv_x + c * shard_col_w
            ax.plot([xx, xx], [kv_y, kv_y + cell_h],
                    color=_TP_LINE_COLOR, linewidth=1.2, linestyle=":")

    ax.set_xlim(-0.3, 4 * (cell_w + gap_x) - gap_x + 0.3)
    ax.set_ylim(-0.9, 2 * (cell_h + gap_y) - gap_y + 0.9)
    ax.set_aspect("auto")
    ax.axis("off")

    ax.set_title(
        f"Tensor sharding view (PE {my_pe})  |  TP={tp}, CP={cp}, "
        f"FFN scope={cfg.topo.ffn_shard_scope}  |  "
        f"blue = this PE's shard",
        fontsize=11, fontweight="bold",
    )
    # Legend for line colors
    ax.plot([], [], color=_TP_LINE_COLOR, linewidth=1.2, linestyle="--",
            label=f"TP split ({tp} ranks)")
    ax.plot([], [], color=_CP_LINE_COLOR, linewidth=1.2, linestyle="--",
            label=f"CP split ({cp} ranks, sequence dim)")
    ax.plot([], [], color=_FFN_LINE_COLOR, linewidth=1.2, linestyle="--",
            label=f"FFN scope split ({cfg.ffn_shard_divisor} ranks)")
    ax.legend(loc="lower center", bbox_to_anchor=(0.5, -0.05),
              ncol=3, fontsize=8, frameon=False)

    return fig
