#!/usr/bin/env python3
"""Conceptual latency-model diagram for the KernBench paper (v5).

Generic naming (Requester Node A/B, Router, Destination Node), with
Router internals visible (two input ports -> switch -> output queue ->
output port) so the queuing delay can be located *at the out port*
rather than at the box edge.  The Destination Node shows queue -> drain
slot -> processing logic.

Highlights:
  * Same-size Routers, each with explicit switching logic + output queue.
  * Edge labels removed -- the wires speak for themselves.
  * Annotation lines are strictly vertical; arrow heads use a larger
    mutation_scale so no thin line protrudes past the tip; `shrinkA`
    pulls the tail away from the text so they no longer overlap.
  * "flit-level interleaving on wires" (not "shared FIFO").
  * Destination: queue -> drain -> processing logic (no "served").

Output:
    docs/report/1H-codesign-paper/figures/latency_model.png

Per /paper isolation: this is a report-only harness under scripts/paper/.
"""

import math
from pathlib import Path

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

OUT = Path(
    "/Users/ywkang/kernbench/docs/report/1H-codesign-paper/figures/latency_model.png"
)
OUT.parent.mkdir(parents=True, exist_ok=True)


# --- Colours --------------------------------------------------------------
C_A      = "#E07B5E"
C_B      = "#5E9BD1"
C_NODE   = "#FFFFFF"
C_BORD   = "#222222"
C_WIRE   = "#222222"
C_ANN    = "#333333"
C_DIM    = "#777777"
C_SWITCH = "#FAFAFA"
C_OQUEUE = "#EFEFEF"
C_PROC   = "#F4ECDC"
C_QUEUE  = "#3CB371"   # medium-sea-green: distinctive vs A/B flit colours


# --- Canvas ---------------------------------------------------------------
fig = plt.figure(figsize=(17.0, 5.0))
ax = fig.add_subplot(111)
ax.set_xlim(0, 33)
ax.set_ylim(2.2, 11.2)
ax.axis("off")


# --- Top header: end-to-end latency formula -----------------------------
ax.text(
    16.5, 10.6,
    r"End-to-end latency  =  $\Sigma$ per-node overhead  +  "
    r"$\Sigma$ per-edge transmission  +  drain  +  "
    r"queuing delay",
    ha="center", fontsize=12, color=C_ANN, weight="bold",
)
ax.plot([1.5, 31.5], [10.05, 10.05], color=C_DIM, lw=0.6)


# --- Helpers --------------------------------------------------------------
def box(cx, cy, w, h, label, fs=11, fweight="normal"):
    ax.add_patch(patches.FancyBboxPatch(
        (cx - w / 2, cy - h / 2), w, h,
        boxstyle="round,pad=0.05",
        facecolor=C_NODE, edgecolor=C_BORD, linewidth=1.6,
    ))
    if label:
        ax.text(cx, cy, label, ha="center", va="center",
                fontsize=fs, weight=fweight)


def draw_flit_aligned(cx, cy, angle_rad, w, h, color, label):
    """Draw a flit polygon rotated to align with the wire angle."""
    cos_a, sin_a = math.cos(angle_rad), math.sin(angle_rad)
    corners = []
    for dx, dy in [(-w / 2, -h / 2), (w / 2, -h / 2),
                   (w / 2, h / 2), (-w / 2, h / 2)]:
        rx = dx * cos_a - dy * sin_a
        ry = dx * sin_a + dy * cos_a
        corners.append((cx + rx, cy + ry))
    ax.add_patch(patches.Polygon(
        corners, facecolor=color, edgecolor="black", linewidth=0.5,
    ))
    ax.text(cx, cy, label, ha="center", va="center",
            fontsize=9, color="white", weight="bold")


def draw_wire_with_flits(x0, y0, x1, y1, n_flits, label_char,
                         color=None, flit_w=0.66, flit_h=0.95, gap=0.10):
    """Wire (line) + endpoint arrowhead + flits centred on the line.

    Wires are drawn in the neutral C_WIRE colour: the wire itself is
    the place where the *transmission* delay accumulates (flit_size /
    BW), not where flits queue. Queuing happens before the wire, at
    the FIFO at the egress side -- coloured separately.
    """
    head_back = 0.30
    dx, dy = x1 - x0, y1 - y0
    wire_len = math.hypot(dx, dy)
    ux, uy = dx / wire_len, dy / wire_len
    x_line_end = x1 - ux * head_back
    y_line_end = y1 - uy * head_back
    ax.plot([x0, x_line_end], [y0, y_line_end],
            color=C_WIRE, lw=1.6, zorder=1)
    head_len, head_half = 0.30, 0.16
    bx = x1 - ux * head_len
    by = y1 - uy * head_len
    px, py = -uy, ux
    tri = [
        (x1, y1),
        (bx + px * head_half, by + py * head_half),
        (bx - px * head_half, by - py * head_half),
    ]
    ax.add_patch(patches.Polygon(tri, facecolor=C_WIRE,
                                  edgecolor=C_WIRE, linewidth=0.0))
    # Flit train
    angle_rad = math.atan2(dy, dx)
    train_len = n_flits * flit_w + (n_flits - 1) * gap
    centre_p = 0.5
    start_p  = centre_p - (train_len / 2) / wire_len
    step_p   = (flit_w + gap) / wire_len
    if isinstance(label_char, str):
        labels = [label_char] * n_flits
        cols   = [color] * n_flits
    else:
        labels = [c[0] for c in label_char]
        cols   = [c[1] for c in label_char]
    for i in range(n_flits):
        p = start_p + (i + 0.5) * step_p
        cx = x0 + p * dx
        cy = y0 + p * dy
        draw_flit_aligned(cx, cy, angle_rad,
                          flit_w, flit_h, cols[i], labels[i])


def draw_router(cx, cy, w, h, two_inputs=True):
    """Same-size Router with explicit internal switching logic and an
    output queue.  Returns the wire-attachment (x,y) for each input
    port and the single output port:
      ((in1_x, in1_y), (in2_x, in2_y) or None, (out_x, out_y))
    """
    # Outer box
    box(cx, cy, w, h, "")
    ax.text(cx, cy + h / 2 - 0.32, "Router",
            ha="center", fontsize=10, weight="bold")

    # Input ports (small circles on the left edge)
    in_x = cx - w / 2
    in_x_internal = in_x + 0.25
    if two_inputs:
        in_y_top = cy + 0.55
        in_y_bot = cy - 0.55
        for iy in (in_y_top, in_y_bot):
            ax.add_patch(patches.Circle(
                (in_x_internal, iy), 0.12,
                facecolor="white", edgecolor=C_BORD, linewidth=0.9,
            ))
    else:
        in_y_top = None
        in_y_bot = cy
        ax.add_patch(patches.Circle(
            (in_x_internal, in_y_bot), 0.12,
            facecolor="white", edgecolor=C_BORD, linewidth=0.9,
        ))

    # Switch (small box, centre-left-ish). sw_cy = cy so that the
    # output queue and (single-input) input port are at the same y as
    # the router centre -- this keeps all router-to-router edges
    # strictly horizontal.
    sw_w, sw_h = 0.85, 1.10
    sw_cx = cx - 0.55
    sw_cy = cy
    # The switch is the router's processing logic -- colour it the
    # same as the Destination Node's processing-logic block so the
    # two read as the same "processing" concept.
    ax.add_patch(patches.Rectangle(
        (sw_cx - sw_w / 2, sw_cy - sw_h / 2), sw_w, sw_h,
        facecolor=C_PROC, edgecolor=C_BORD, linewidth=0.8,
    ))
    ax.text(sw_cx, sw_cy, "switch", ha="center", va="center",
            fontsize=7.5, style="italic")

    # Short feeder lines (no arrowhead) from input ports to the
    # switch.  We deliberately drop arrowheads here: the head + port
    # circle were too small at this scale and read as an overlap.
    sw_left_x = sw_cx - sw_w / 2
    if two_inputs:
        for iy in (in_y_top, in_y_bot):
            ax.plot(
                [in_x_internal + 0.12, sw_left_x],
                [iy, sw_cy + 0.25 * ((iy - cy) / 0.55)],
                color=C_DIM, lw=0.7, zorder=1,
            )
    else:
        ax.plot(
            [in_x_internal + 0.12, sw_left_x],
            [in_y_bot, sw_cy],
            color=C_DIM, lw=0.7, zorder=1,
        )

    # Output queue (small queue holding a couple of flits). This *is*
    # a real queueing location, so its fill takes the C_QUEUE family.
    oq_w, oq_h = 0.95, 0.55
    oq_cx = cx + 0.55
    oq_cy = sw_cy
    ax.add_patch(patches.Rectangle(
        (oq_cx - oq_w / 2, oq_cy - oq_h / 2), oq_w, oq_h,
        facecolor="#D9F0E1", edgecolor=C_QUEUE, linewidth=0.9,
    ))
    ax.text(oq_cx, oq_cy + oq_h / 2 + 0.18, "FIFO",
            ha="center", fontsize=7, color=C_DIM, style="italic")
    # Two small flits inside (A and B) hinting at the in-flight contents
    mini_w, mini_h = 0.22, 0.34
    for i, (col, lab) in enumerate([(C_A, "A"), (C_B, "B")]):
        mx = oq_cx - 0.30 + i * (mini_w + 0.06)
        ax.add_patch(patches.Rectangle(
            (mx, oq_cy - mini_h / 2), mini_w, mini_h,
            facecolor=col, edgecolor="black", linewidth=0.3,
        ))
        ax.text(mx + mini_w / 2, oq_cy, lab,
                ha="center", va="center",
                fontsize=5.5, color="white", weight="bold")

    # Short feeder line (no arrowhead) from switch into out queue
    ax.plot(
        [sw_cx + sw_w / 2, oq_cx - oq_w / 2],
        [sw_cy, oq_cy],
        color=C_DIM, lw=0.7, zorder=1,
    )
    # Short feeder line from out queue to out port
    out_x_internal = cx + w / 2 - 0.25
    ax.plot(
        [oq_cx + oq_w / 2, out_x_internal - 0.12],
        [oq_cy, oq_cy],
        color=C_DIM, lw=0.7, zorder=1,
    )
    # Output port circle on the right edge (a port marker, not a
    # queue -- the queue is on the wire that follows, not the port).
    ax.add_patch(patches.Circle(
        (out_x_internal, oq_cy), 0.12,
        facecolor="white", edgecolor=C_BORD, linewidth=0.9,
    ))

    return (
        (in_x_internal, in_y_top) if two_inputs else None,
        (in_x_internal, in_y_bot),
        (out_x_internal, oq_cy),
        oq_cx,   # also return the out-queue centre so callouts can target it
        oq_cy,
    )


# --- Box layout ---------------------------------------------------------
# Requester centres set to match Router 1's input-port y so that
# Edge 1A and Edge 1B are strictly horizontal. Box heights are kept
# small enough that a visible gap separates the two Requester boxes.
R1_TMP_CY = 7.0
ReqA = (2.8, R1_TMP_CY + 0.55)   # = 7.55, matches in_y_top of R1
ReqB = (2.8, R1_TMP_CY - 0.55)   # = 6.45, matches in_y_bot of R1
box(*ReqA, 3.0, 0.85, "Requester\nNode A", fs=10)
box(*ReqB, 3.0, 0.85, "Requester\nNode B", fs=10)

R_W, R_H = 3.0, 2.8
R1 = (10.5, 7.0)
R2 = (20.0, 7.0)

r1_in_top, r1_in_bot, r1_out, r1_oq_cx, r1_oq_cy = draw_router(
    *R1, R_W, R_H, two_inputs=True,
)
_, r2_in_bot, r2_out, r2_oq_cx, r2_oq_cy = draw_router(
    *R2, R_W, R_H, two_inputs=False,
)

# Destination Node (same height as Router; wide enough so queue +
# drain + processing-logic all fit on a single horizontal row).
Dst = (28.4, 7.0)
Dst_W, Dst_H = 6.0, R_H   # match the Router height
box(*Dst, Dst_W, Dst_H, "")
ax.text(Dst[0], Dst[1] + Dst_H / 2 - 0.25, "Destination Node",
        ha="center", fontsize=10, weight="bold")


# --- Edges: requester -> router 1 (Edge 1A & 1B, with flits) ------------
draw_wire_with_flits(
    ReqA[0] + 1.6, ReqA[1],
    r1_in_top[0] - 0.02, r1_in_top[1],
    n_flits=4, color=C_A, label_char="A",
)
draw_wire_with_flits(
    ReqB[0] + 1.6, ReqB[1],
    r1_in_bot[0] - 0.02, r1_in_bot[1],
    n_flits=4, color=C_B, label_char="B",
)

# Edge 2: router1 out -> router2 in (horizontal, interleaved)
labels_e2 = [("A", C_A) if i % 2 == 0 else ("B", C_B) for i in range(8)]
E2_y = r1_out[1]
draw_wire_with_flits(
    r1_out[0] + 0.02, r1_out[1],
    r2_in_bot[0] - 0.02, r2_in_bot[1],
    n_flits=8, label_char=labels_e2,
)

# Edge 3: router2 out -> destination (horizontal, interleaved)
labels_e3 = [("A", C_A) if i % 2 == 0 else ("B", C_B) for i in range(4)]
draw_wire_with_flits(
    r2_out[0] + 0.02, r2_out[1],
    Dst[0] - Dst_W / 2, r2_out[1],
    n_flits=4, label_char=labels_e3,
)


# --- Destination internals: queue -> drain -> processing logic ----------
# Light-green halo behind the queue area marks it as a queueing point.
C_QUEUE_BG = "#D9F0E1"

dy_main = Dst[1] - 0.10
Dst_left   = Dst[0] - Dst_W / 2

# Queue (4 flits, FIFO; rightmost is the next to drain, so the
# order is set so the serve sequence after the in-flight A alternates
# A (in drain) -> B -> A -> B -> A -- giving a clean BABA queue.
qW, qH, qGap = 0.42, 0.62, 0.07
q_labels = ["A", "B", "A", "B"]
q_cols   = [C_A,  C_B,  C_A,  C_B]
q_total  = len(q_labels) * qW + (len(q_labels) - 1) * qGap
q_x_start = Dst_left + 0.35

# Green halo behind the queue boxes (the queue itself is a queueing
# location -- mark it with the C_QUEUE colour family)
ax.add_patch(patches.FancyBboxPatch(
    (q_x_start - 0.10, dy_main - qH / 2 - 0.08),
    q_total + 0.20, qH + 0.16,
    boxstyle="round,pad=0.02",
    facecolor=C_QUEUE_BG, edgecolor=C_QUEUE,
    linewidth=0.9, zorder=1.5,
))

for i, (lab, col) in enumerate(zip(q_labels, q_cols)):
    qx = q_x_start + i * (qW + qGap)
    ax.add_patch(patches.Rectangle(
        (qx, dy_main - qH / 2), qW, qH,
        facecolor=col, edgecolor="black", linewidth=0.4,
        zorder=2,
    ))
    ax.text(qx + qW / 2, dy_main, lab,
            ha="center", va="center",
            fontsize=8, color="white", weight="bold", zorder=3)
# Common y for the "queue" / "drain" labels. Matches the FIFO label
# spacing inside the Router (0.18 above the box top) and uses the same
# font size for visual consistency.
DST_LABEL_Y = dy_main + qH / 2 + 0.18
ax.text(q_x_start + q_total / 2, DST_LABEL_Y,
        "queue",
        ha="center", fontsize=7, style="italic", color=C_DIM)

# Drain slot (height matched to queue boxes; font matched to FIFO/queue
# labels for visual consistency)
dr_W, dr_H = 0.85, qH
dr_x = q_x_start + q_total + 0.35
dr_cx = dr_x + dr_W / 2
dr_cy = dy_main
ax.add_patch(patches.FancyBboxPatch(
    (dr_x, dr_cy - dr_H / 2), dr_W, dr_H,
    boxstyle="round,pad=0.03",
    facecolor=C_A, edgecolor="black", linewidth=1.0,
))
ax.text(dr_cx, dr_cy, "A", ha="center", va="center",
        fontsize=9, color="white", weight="bold")
ax.text(dr_cx, DST_LABEL_Y,
        "drain",
        ha="center", fontsize=7, style="italic", color=C_DIM)

# Drain-time double-arrow underneath the drain slot. Move the "drain
# time" text a touch further down so the descender does not collide
# with the arrowhead.
dt_y = dr_cy - dr_H / 2 - 0.28
ax.annotate(
    "", xy=(dr_x + dr_W, dt_y), xytext=(dr_x, dt_y),
    arrowprops=dict(arrowstyle="<->", lw=0.9, color=C_ANN),
)
ax.text(dr_cx, dt_y - 0.42, "drain time",
        ha="center", fontsize=8.5, color=C_ANN, style="italic")

# Processing-logic block (after drain) — wider so the two-line text
# does not collide with the box edges
pl_W, pl_H = 1.75, 1.00
pl_x = dr_x + dr_W + 0.40
pl_cx = pl_x + pl_W / 2
pl_cy = dr_cy
ax.add_patch(patches.FancyBboxPatch(
    (pl_x, pl_cy - pl_H / 2), pl_W, pl_H,
    boxstyle="round,pad=0.03",
    facecolor=C_PROC, edgecolor=C_BORD, linewidth=1.0,
))
ax.text(pl_cx, pl_cy, "processing\nlogic",
        ha="center", va="center", fontsize=9)

# Small gray arrows along the queue -> drain -> processing chain
ax.annotate(
    "", xy=(dr_x - 0.04, dr_cy),
    xytext=(q_x_start + q_total + 0.13, dr_cy),
    arrowprops=dict(arrowstyle="-|>", lw=0.7,
                    color=C_DIM, mutation_scale=7),
)
ax.annotate(
    "", xy=(pl_x - 0.04, dr_cy),
    xytext=(dr_x + dr_W + 0.04, dr_cy),
    arrowprops=dict(arrowstyle="-|>", lw=0.7,
                    color=C_DIM, mutation_scale=7),
)


# --- Annotations (strictly vertical, no diagonals, no overlap) ---------

def vcallout(text, xy, x_text_top_y, fs=10, clearance=0.55,
             marker_color=None):
    """Vertical dotted callout. Optional `marker_color` paints a small
    coloured circle just left of the text -- used to tie the label to
    a colour code in the diagram.
    """
    text_y = x_text_top_y
    if text_y < xy[1]:
        line_top_y = text_y + clearance
    else:
        line_top_y = text_y - clearance
    ax.plot([xy[0], xy[0]], [xy[1], line_top_y],
            color=C_ANN, lw=0.9, linestyle=":", zorder=2)
    if marker_color is not None:
        # Text-width estimate at fs=10 (~0.18 data units per char) so
        # the marker is placed clearly to the left of the text.
        text_w = len(text) * 0.18
        marker_x = xy[0] - text_w / 2 - 0.30
        ax.add_patch(patches.Circle(
            (marker_x, text_y), 0.14,
            facecolor=marker_color, edgecolor=C_BORD, linewidth=0.6,
            zorder=3,
        ))
    ax.text(xy[0], text_y, text,
            ha="center", va="center",
            fontsize=fs, color=C_ANN, zorder=3)


# Transmission delay -> flit on Edge 2 (text BELOW the wire).
# Extra clearance ~ 2 x text height so the line stops well clear of
# the label.
e2_total = 8 * 0.66 + 7 * 0.10
e2_centre = (r1_out[0] + r2_in_bot[0]) / 2
e2_start = e2_centre - e2_total / 2
trans_idx = 5
trans_cx  = e2_start + (trans_idx + 0.5) * (0.66 + 0.10)
vcallout(
    "transmission delay  =  flit_size / BW",
    xy=(trans_cx, E2_y - 0.55),
    x_text_top_y=E2_y - 2.6,
    clearance=0.45,
)

# Flit-level interleaving -> ABOVE the wire, anchored on a flit near
# the centre/right of the wire. Text y matches the queuing-delay
# callout so the two labels sit on the same horizontal row, separated
# horizontally so they don't collide.
int_idx = 4
int_cx  = e2_start + (int_idx + 0.5) * (0.66 + 0.10)
vcallout(
    "flit-level interleaving on wires",
    xy=(int_cx, E2_y + 0.55),
    x_text_top_y=R1[1] + R_H / 2 + 0.85,
)

# Queuing delay -> at Router 1 out-queue. Text positioned just above
# the Router-1 box.
vcallout(
    "queuing delay",
    xy=(r1_oq_cx, r1_oq_cy + 0.30),
    x_text_top_y=R1[1] + R_H / 2 + 0.85,
    marker_color=C_QUEUE,
)

# Drain -> below the drain slot (text is below).
# Extra clearance ~ 2.5 x text height so the line stops further from
# the label.
vcallout(
    "drain  =  per-flit service occupancy",
    xy=(dr_cx, dt_y - 0.66),   # just *inside* the Destination Node
                                  # bottom edge -- the dotted line then
                                  # penetrates the box slightly rather
                                  # than hanging below it
    x_text_top_y=E2_y - 2.6,
    clearance=0.60,
)

# Per-node overhead -> points at Router 1's switch block (where the
# component's fixed processing cost lives). Text on the same row as
# transmission delay and drain so the three "below the wire" callouts
# sit on one horizontal baseline.
R1_SW_CX = R1[0] - 0.55         # sw_cx for Router 1
R1_SW_BOTTOM_Y = R1[1] - 0.55   # bottom of sw box
vcallout(
    "per-node overhead",
    xy=(R1_SW_CX, R1_SW_BOTTOM_Y),
    x_text_top_y=E2_y - 2.6,
    clearance=0.60,
    marker_color=C_PROC,
)


# --- Legend (close to the diagram) --------------------------------------
LX, LY = 1.0, 2.6
ax.add_patch(patches.Rectangle((LX, LY), 0.7, 0.45,
                                facecolor=C_A, edgecolor="black",
                                linewidth=0.4))
ax.text(LX + 0.95, LY + 0.22, "Transaction A flit",
        va="center", fontsize=10)

ax.add_patch(patches.Rectangle((LX + 4.8, LY), 0.7, 0.45,
                                facecolor=C_B, edgecolor="black",
                                linewidth=0.4))
ax.text(LX + 5.75, LY + 0.22, "Transaction B flit",
        va="center", fontsize=10)


# --- Save ----------------------------------------------------------------
fig.savefig(OUT, dpi=140, bbox_inches="tight", facecolor="white")
print(f"Wrote {OUT}")