"""Generate 3 perspective (isometric) SVGs for the MHSA attention story:

    06_attn_raw.svg     —  raw attention pattern, no mask
    07_attn_mask.svg    —  the causal-packet mask itself (binary shape)
    08_attn_masked.svg  —  raw attention with the mask applied

Each cell is rendered as a parallelogram via a 30°/30° isometric projection.
All three SVGs share the same dimensions and the same projection so they can
be composed / cross-faded in drawio.
"""
from pathlib import Path
import math

# --- palette ----------------------------------------------------------------
CORANGE = (230, 126, 34)
WHITE   = (255, 255, 255)

# --- geometry ---------------------------------------------------------------
CELL = 28           # local cell size (before projection)
N    = 12           # 12 x 12 grid
PAD  = 24

COS30 = math.sqrt(3) / 2     # ≈ 0.866
SIN30 = 0.5


def project(x: float, y: float) -> tuple[float, float]:
    """Map local (x, y) to screen (sx, sy) via 30°/30° isometric projection.

    Local x-axis (cols) goes down-right, local y-axis (rows) goes down-left.
    Origin (0, 0) ends up at the TOP of the projected diamond.
    """
    sx = x * COS30 - y * COS30
    sy = x * SIN30 + y * SIN30
    return sx, sy


def _bbox() -> tuple[float, float, float, float]:
    corners = [
        project(0, 0),
        project(N * CELL, 0),
        project(0, N * CELL),
        project(N * CELL, N * CELL),
    ]
    xs, ys = [c[0] for c in corners], [c[1] for c in corners]
    return min(xs), max(xs), min(ys), max(ys)


_xmin, _xmax, _ymin, _ymax = _bbox()
W = int(_xmax - _xmin + 2 * PAD)
H = int(_ymax - _ymin + 2 * PAD)
OX = -_xmin + PAD
OY = -_ymin + PAD


def mix_orange(p: int) -> str:
    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_poly(i: int, j: int) -> str:
    """Return SVG points string for the parallelogram at row i, col j."""
    pts = []
    for di, dj in [(0, 0), (0, 1), (1, 1), (1, 0)]:
        x = (j + dj) * CELL
        y = (i + di) * CELL
        sx, sy = project(x, y)
        pts.append(f"{sx + OX:.2f},{sy + OY:.2f}")
    return " ".join(pts)


def make_svg(style_fn) -> str:
    # No background rect — SVGs are transparent so they can be cleanly
    # overlaid / placed on any drawio canvas color.
    parts = [
        f'<svg xmlns="http://www.w3.org/2000/svg" '
        f'width="{W}" height="{H}" viewBox="0 0 {W} {H}">',
    ]
    # draw row-by-row, far rows first so near rows overlay (depth order).
    for i in range(N):
        for j in range(N):
            fill, stroke, sw = style_fn(i, j)
            parts.append(
                f'  <polygon points="{cell_poly(i, j)}" '
                f'fill="{fill}" stroke="{stroke}" stroke-width="{sw}"/>'
            )
    parts.append('</svg>')
    return '\n'.join(parts) + '\n'


# ----------------------------------------------------------------------------
# Style functions for each of the 3 SVGs
# ----------------------------------------------------------------------------

def style_raw(i: int, j: int) -> tuple[str, str, float]:
    """Raw attention before any mask: band-matrix-like soft pattern.

    Brightness peaks on the diagonal and decays with |i - j|, giving a
    realistic-looking dense attention map.
    """
    dist = abs(i - j)
    opa = max(20, 75 - dist * 4)
    return mix_orange(opa), "#9C9C9C", 0.4


def style_mask(i: int, j: int) -> tuple[str, str, float]:
    """The causal-packet mask itself, as a binary visual:

        - allowed cells (FLOW row/col + lower-tri):   orange-light
        - blocked cells (upper-tri except FLOW):      dark gray
    """
    blocked = (i != 0) and (j != 0) and (j > i)
    if blocked:
        return "#3A3A3A", "#1F1F1F", 0.5
    return mix_orange(60), "#B8651F", 0.4


def style_masked(i: int, j: int) -> tuple[str, str, float]:
    """Raw attention AFTER applying the causal-packet mask: blocked cells
    are whited out, allowed cells keep EXACTLY their raw attention intensity
    (i.e. masked == raw × mask, nothing else changes).
    """
    blocked = (i != 0) and (j != 0) and (j > i)
    if blocked:
        return "#FFFFFF", "#CCCCCC", 0.4
    return style_raw(i, j)


OUT = Path(__file__).parent
(OUT / "06_attn_raw.svg").write_text(make_svg(style_raw))
(OUT / "07_attn_mask.svg").write_text(make_svg(style_mask))
(OUT / "08_attn_masked.svg").write_text(make_svg(style_masked))
print(f"wrote 06_attn_raw.svg / 07_attn_mask.svg / 08_attn_masked.svg ({W}x{H})")
