"""Pipeline diagram: attention + FFN steps, colored by bound type.

Each stage is drawn as a rounded box in a horizontal flow:
  Attention: S1 -> S2 -> S3 -> ... -> S10 -> C2 (AllReduce)
  FFN:       F1 -> F2 -> F3 -> F4 -> F5 -> CF1 (AllReduce)

CP ring communication (C1) and score AllReduce (C3) are drawn as
"overhead" boxes attached above the QKt/PV region (they run
concurrently with per-hop compute in the ring).

Box fill color = the stage's bound classification:
  compute-bound  blue
  memory-bound   orange
  comm-bound     red
  trivial        grey
"""
from __future__ import annotations

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

from .model_config import FullConfig
from .stage_latencies import (
    all_stages, all_ffn_stages, StageCost,
)


BOUND_COLOR = {
    "compute": "#4682b4",   # blue
    "memory":  "#e37400",   # orange
    "comm":    "#c62828",   # red
    "trivial": "#a9a9a9",   # grey
}
BOUND_LABEL = {
    "compute": "compute-bound",
    "memory":  "memory-bound",
    "comm":    "communication",
    "trivial": "trivial / negligible",
}


def _fmt_us(sec: float) -> str:
    us = sec * 1e6
    if us < 1:
        return f"{us*1e3:.1f}ns"
    if us < 100:
        return f"{us:.2f}us"
    return f"{us:.0f}us"


def _draw_block(ax, x, y, w, h, label, cost: StageCost, edge="#333"):
    color = BOUND_COLOR.get(cost.bound, "#a9a9a9")
    box = patches.FancyBboxPatch(
        (x, y), w, h,
        boxstyle="round,pad=0.05",
        facecolor=color, edgecolor=edge, linewidth=1.2, alpha=0.85,
    )
    ax.add_patch(box)
    # Two lines: stage name (bold) + visible latency (small)
    ax.text(x + w / 2, y + h * 0.62, label,
            ha="center", va="center",
            fontsize=8.0, fontweight="bold", color="white")
    ax.text(x + w / 2, y + h * 0.25, _fmt_us(cost.visible_s),
            ha="center", va="center",
            fontsize=7.0, color="white")


_SHORT_NAMES = {
    "S1":  "RMSNorm",
    "S2":  "W_Q",
    "S3":  "W_K+W_V",
    "S4":  "KV append",
    "S5":  "Q.K^T",
    "S6":  "softmax",
    "S7":  "P.V",
    "S8":  "merge",
    "S9":  "norm O/l",
    "S10": "W_O",
    "C1":  "CP ring",
    "C2":  "TP AR",
    "C3":  "Score AR",
    "F1":  "RMSNorm",
    "F2":  "W_gate",
    "F3":  "W_up",
    "F4":  "SwiGLU",
    "F5":  "W_down",
    "CF1": "FFN AR",
}


def _short_name(name: str) -> str:
    """Map 'S1 RMSNorm' -> 'RMSNorm'; fallback to first token."""
    tok = name.split()[0]
    return _SHORT_NAMES.get(tok, tok)


def _draw_arrow(ax, x0, x1, y):
    ax.annotate("", xy=(x1, y), xytext=(x0, y),
                arrowprops=dict(arrowstyle="->", color="#333", lw=1.1))


def _draw_ring_loop(ax, x_left, x_right, y_top, label: str, n_iters: int,
                     arc_lift: float = 0.9):
    """Draw a curved back-arrow above [x_left..x_right] with a repeat label.

    Visualizes 'repeat this block n_iters times' (like a loop in a diagram).
    arc_lift: how high above y_top the arc sits.
    """
    color = "#5a189a"
    lw = 1.6
    arc_y = y_top + arc_lift
    # Vertical tick up from the right edge
    ax.plot([x_right, x_right], [y_top + 0.02, arc_y],
            color=color, linewidth=lw)
    # Horizontal top of the loop
    ax.plot([x_right, x_left], [arc_y, arc_y],
            color=color, linewidth=lw)
    # Down-arrow into the left edge (a back-edge indicating the repeat)
    ax.annotate("", xy=(x_left, y_top + 0.02), xytext=(x_left, arc_y),
                arrowprops=dict(arrowstyle="->", color=color, lw=lw))
    # Loop label centered above the arc
    ax.text((x_left + x_right) / 2, arc_y + 0.08,
            label,
            ha="center", va="bottom",
            fontsize=9, fontweight="bold", color=color)
    # Iteration count as a badge on the right
    ax.text(x_right + 0.08, arc_y - 0.03,
            f"x {n_iters}",
            ha="left", va="top",
            fontsize=10, fontweight="bold", color=color,
            bbox=dict(boxstyle="round,pad=0.15", facecolor="white",
                      edgecolor=color, linewidth=1.2, alpha=0.95))


def _draw_row(ax, y_base, row_title, stages: list[StageCost],
              x_start: float, box_w: float, box_h: float, gap: float):
    """Draw one row of stage boxes with arrows between them."""
    ax.text(x_start - 0.4, y_base + box_h / 2, row_title,
            ha="right", va="center",
            fontsize=11, fontweight="bold", color="#212529")
    x = x_start
    for i, st in enumerate(stages):
        _draw_block(ax, x, y_base, box_w, box_h,
                    _short_name(st.name), st)
        if i < len(stages) - 1:
            _draw_arrow(ax, x + box_w + 0.02, x + box_w + gap - 0.02,
                        y_base + box_h / 2)
        x += box_w + gap
    return x - gap


def _draw_overlap_boxes(ax, hidden_stages, x_left, x_right, y_top,
                          gap_between: float = 0.15):
    """Draw dashed 'overlapped comm' boxes spanning x_left..x_right above y_top.

    Boxes are widened to cover the horizontal range of the stages they
    overlap with (e.g. C1 CP ring sits above S5-S7).
    """
    if not hidden_stages:
        return
    hy = y_top + 0.10
    n = len(hidden_stages)
    total_span = x_right - x_left
    box_span = (total_span - (n - 1) * gap_between) / max(1, n)
    box_h_local = 0.42
    hx = x_left
    for st in hidden_stages:
        box = patches.FancyBboxPatch(
            (hx, hy), box_span, box_h_local,
            boxstyle="round,pad=0.03",
            facecolor=BOUND_COLOR["comm"], edgecolor="#333",
            linewidth=1.0, alpha=0.35, linestyle="--",
        )
        ax.add_patch(box)
        ax.text(hx + box_span / 2, hy + box_h_local * 0.65,
                f"{_short_name(st.name)} (concurrent)",
                ha="center", va="center",
                fontsize=7, color="#7a0e0e", fontweight="bold")
        ax.text(hx + box_span / 2, hy + box_h_local * 0.2,
                _fmt_us(st.visible_s),
                ha="center", va="center",
                fontsize=6.5, color="#7a0e0e")
        hx += box_span + gap_between


def draw_pipeline(cfg: FullConfig, ax=None):
    """Draw attention + FFN pipeline. Boxes colored by bound type."""
    attn = all_stages(cfg)
    ffn = all_ffn_stages(cfg)

    # Split attention: main stages, then C2 (TP AllReduce). C1/C3 = overlapped.
    attn_main = [s for s in attn if not s.name.startswith("C")]
    attn_c2 = next((s for s in attn if s.name.startswith("C2")), None)
    attn_hidden = [s for s in attn
                   if s.name.startswith("C1") or s.name.startswith("C3")]
    attn_hidden = [s for s in attn_hidden if s.visible_s > 0]
    attn_row = attn_main + ([attn_c2] if attn_c2 and attn_c2.visible_s > 0 else [])

    # FFN row: all F* + CF1 if non-trivial
    ffn_main = [s for s in ffn if not s.name.startswith("CF")]
    ffn_cf = next((s for s in ffn if s.name.startswith("CF1")), None)
    ffn_row = ffn_main + ([ffn_cf] if ffn_cf and ffn_cf.visible_s > 0 else [])

    n_max = max(len(attn_row), len(ffn_row))
    box_w = 1.35
    box_h = 0.78
    gap = 0.18
    x_start = 1.1

    total_w = x_start + n_max * box_w + (n_max - 1) * gap + 0.5
    total_h = 4.2

    if ax is None:
        fig, ax = plt.subplots(figsize=(max(11, total_w * 1.1),
                                         max(4.5, total_h * 1.0)))
    else:
        fig = ax.figure

    _draw_row(ax, y_base=2.6, row_title="Attention",
              stages=attn_row, x_start=x_start,
              box_w=box_w, box_h=box_h, gap=gap)

    _draw_row(ax, y_base=0.6, row_title="FFN",
              stages=ffn_row, x_start=x_start,
              box_w=box_w, box_h=box_h, gap=gap)

    # Ring loop indicator + concurrent-comm boxes anchored over S5..S8.
    _idx_by_prefix = {s.name.split()[0]: i for i, s in enumerate(attn_row)}
    _s5_idx = _idx_by_prefix.get("S5")
    _s7_idx = _idx_by_prefix.get("S7")
    _s8_idx = _idx_by_prefix.get("S8")

    if cfg.topo.cp > 1 and cfg.topo.mode == "prefill" and _s5_idx is not None:
        _right_idx = _s8_idx if _s8_idx is not None else _s7_idx
        _lx = x_start + _s5_idx * (box_w + gap)
        _rx = x_start + _right_idx * (box_w + gap) + box_w
        _y_top = 2.6 + box_h
        # Concurrent-comm boxes (C1 CP ring, C3 score AR) directly above S5-S7
        _draw_overlap_boxes(ax, attn_hidden, _lx, _rx, _y_top)
        # Ring loop arc above those
        _draw_ring_loop(
            ax, _lx, _rx, _y_top,
            label="ring attention loop (K/V ring per hop; S8 merges partial O,m,l)",
            n_iters=cfg.topo.cp,
            arc_lift=1.05,
        )
    elif cfg.topo.cp > 1 and cfg.topo.mode == "decode" and _s5_idx is not None:
        _right_idx = _s7_idx if _s7_idx is not None else _s5_idx
        _lx = x_start + _s5_idx * (box_w + gap)
        _rx = x_start + _right_idx * (box_w + gap) + box_w
        _y_top = 2.6 + box_h
        # Any straggling concurrent comm (e.g. C3 if head-split active) above S5-S7
        _draw_overlap_boxes(ax, attn_hidden, _lx, _rx, _y_top)
        _draw_ring_loop(
            ax, _lx, _rx, _y_top,
            label="decode: local pass only; S8 fires O/m/l all-reduce",
            n_iters=1,
            arc_lift=1.05 if attn_hidden else 0.85,
        )
    elif attn_hidden:
        # Fallback: CP=1 but still have C3 or similar; anchor above S5-S7 if present
        if _s5_idx is not None and _s7_idx is not None:
            _lx = x_start + _s5_idx * (box_w + gap)
            _rx = x_start + _s7_idx * (box_w + gap) + box_w
            _draw_overlap_boxes(ax, attn_hidden, _lx, _rx, 2.6 + box_h)

    # Totals
    attn_total = sum(s.visible_s for s in attn_row) \
        + sum(s.visible_s for s in attn_hidden)
    ffn_total = sum(s.visible_s for s in ffn_row)
    ax.text(total_w - 0.3, 2.6 + box_h / 2 + 1.15,
            f"Attn total: {_fmt_us(attn_total)}",
            ha="right", va="bottom", fontsize=9,
            fontweight="bold", color="#212529")
    ax.text(total_w - 0.3, 0.6 + box_h + 0.15,
            f"FFN total:  {_fmt_us(ffn_total)}",
            ha="right", va="bottom", fontsize=9,
            fontweight="bold", color="#212529")

    # Legend
    handles = []
    for k, lbl in BOUND_LABEL.items():
        handles.append(patches.Patch(
            facecolor=BOUND_COLOR[k], edgecolor="#333",
            alpha=0.85, label=lbl,
        ))
    ax.legend(handles=handles, loc="upper center",
              bbox_to_anchor=(0.5, -0.02),
              ncol=4, fontsize=8, frameon=False)

    ax.set_xlim(0, total_w + 0.6)
    ax.set_ylim(0, 4.9)
    ax.set_aspect("auto")
    ax.axis("off")
    ax.set_title(
        f"Per-layer pipeline (one PE, {cfg.topo.mode}, T_q={cfg.topo.T_q}, "
        f"S_local={cfg.topo.s_local}) - color = dominant bound",
        fontsize=10, fontweight="bold",
    )
    return fig