"""Generate the causal-packet MHSA attention heatmap SVG.

12 x 12 attention matrix, rounded-rectangle cells, no text labels:
    - row 0 and col 0 (FLOW token):       cOrange!85  (full attention)
    - lower triangular within [1:, 1:]:   cOrange decaying with distance
                                          (cOrange!{75 - 4*dist})
    - upper triangular (blocked):         white + light grey border

Run:
    python gen_attn_heatmap.py
"""
from pathlib import Path

CORANGE = (230, 126, 34)        # RGB of cOrange palette
WHITE   = (255, 255, 255)

CELL    = 30                    # px per cell (square)
GAP     = 2                     # px between cells
PAD     = 6                     # px outer padding
RADIUS  = 4                     # px corner radius for cells
N       = 12                    # 12 x 12 grid

STRIDE  = CELL + GAP
SIZE    = N * CELL + (N - 1) * GAP + 2 * PAD


def mix_orange(p: int) -> str:
    """cOrange mixed with white at p% (matches TikZ cOrange!p)."""
    f = p / 100
    r = round(CORANGE[0] * f + WHITE[0] * (1 - f))
    g = round(CORANGE[1] * f + WHITE[1] * (1 - f))
    b = round(CORANGE[2] * f + WHITE[2] * (1 - f))
    return f"#{r:02X}{g:02X}{b:02X}"


def cell_style(i: int, j: int) -> tuple[str, str, float]:
    """(fill, stroke, stroke_width) for cell (row=i, col=j)."""
    if i == 0 or j == 0:                                # FLOW row / col
        return mix_orange(85), "#B8651F", 0.6
    if j > i:                                            # blocked (upper-tri)
        return "#FFFFFF", "#BBBBBB", 0.5
    opa = max(15, 75 - (i - j) * 4)                      # lower-tri decay
    return mix_orange(opa), "#999999", 0.3


def make_svg() -> str:
    parts = [
        f'<svg xmlns="http://www.w3.org/2000/svg" '
        f'width="{SIZE}" height="{SIZE}" viewBox="0 0 {SIZE} {SIZE}">',
        f'  <rect width="{SIZE}" height="{SIZE}" fill="white"/>',
    ]
    for i in range(N):
        for j in range(N):
            x = PAD + j * STRIDE
            y = PAD + i * STRIDE
            fill, stroke, sw = cell_style(i, j)
            parts.append(
                f'  <rect x="{x}" y="{y}" width="{CELL}" height="{CELL}" '
                f'rx="{RADIUS}" ry="{RADIUS}" '
                f'fill="{fill}" stroke="{stroke}" stroke-width="{sw}"/>'
            )
    parts.append('</svg>')
    return '\n'.join(parts) + '\n'


OUT = Path(__file__).parent
(OUT / "05_attention_heatmap.svg").write_text(make_svg())
print(f"wrote 05_attention_heatmap.svg ({SIZE}x{SIZE})")
