"""Generate 4 SVG tensor visualisations for the AdaLN modulation flow.

Each SVG is a 6 col x 4 row grid where cell color = cBlue mixed with white at
the per-cell opacity (0..100). Same color values used in the TikZ panel (c).

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

CBLUE = (52, 110, 180)         # RGB of cBlue palette in the TikZ figure
WHITE = (255, 255, 255)
CELL = 50                       # px per cell
BORDER = 3                      # outer padding (px)
COLS, ROWS = 6, 4
W = COLS * CELL + 2 * BORDER
H = ROWS * CELL + 2 * BORDER


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


def make_svg(opacities: list[int]) -> str:
    parts = [
        f'<svg xmlns="http://www.w3.org/2000/svg" '
        f'width="{W}" height="{H}" viewBox="0 0 {W} {H}">',
        f'  <rect width="{W}" height="{H}" fill="white"/>',
    ]
    for k, op in enumerate(opacities):
        i, j = k % COLS, k // COLS
        x = BORDER + i * CELL
        y = BORDER + j * CELL
        parts.append(
            f'  <rect x="{x}" y="{y}" width="{CELL}" height="{CELL}" '
            f'fill="{mix(op)}" stroke="#999" stroke-width="0.6"/>'
        )
    parts.append('</svg>')
    return '\n'.join(parts) + '\n'


# Same opacity arrays used in the TikZ panel (c).  Row-major: 6 cols * 4 rows.
TENSORS = {
    "01_h_in":         [30, 80, 50, 40, 90, 60,
                        70, 20, 80, 50, 40, 90,
                        60, 50, 90, 30, 80, 40,
                        70, 30, 60, 80, 40, 70],
    "02_h_after_LN":   [50, 30, 70, 50, 40, 60,
                        40, 60, 50, 50, 40, 60,
                        50, 50, 60, 40, 60, 40,
                        50, 40, 50, 60, 40, 60],
    "03_h_after_gamma":[60, 12, 70, 25, 52, 60,
                        48, 24, 50, 25, 52, 60,
                        60, 20, 60, 20, 78, 40,
                        60, 16, 50, 30, 52, 60],
    "04_h_after_beta": [80, 32, 90, 45, 72, 80,
                        68, 44, 70, 45, 72, 80,
                        80, 40, 80, 40, 90, 60,
                        80, 36, 70, 50, 72, 80],
}

OUT = Path(__file__).parent
for name, vals in TENSORS.items():
    assert len(vals) == COLS * ROWS, f"{name}: need {COLS*ROWS} values"
    (OUT / f"{name}.svg").write_text(make_svg(vals))
    print(f"wrote {name}.svg ({W}x{H})")
