#!/usr/bin/env python3
"""Render the forward-pass figure to PNG. Layout mirrors forward-pass-on-the-graph.svg."""
from PIL import Image, ImageDraw, ImageFont

S = 2  # supersample scale
W, H = 1320 * S, 640 * S

INK      = (36, 48, 63)
MUTED    = (90, 104, 117)
ACCENT   = (11, 107, 98)
GREEN    = (47, 125, 95)
RED      = (194, 84, 58)
FRAME    = (201, 210, 221)
GRAY     = (138, 152, 168)
PALE     = (185, 195, 206)

FILL_STAGE  = (245, 247, 250)
FILL_TERM   = (238, 247, 242)
FILL_REFUSE = (253, 241, 237)
FILL_STORE  = (251, 252, 253)

FONT_DIR = "/System/Library/Fonts/Supplemental/"
def f(name, px):
    return ImageFont.truetype(FONT_DIR + name, int(px * S))

F_LABEL = f("Arial Bold.ttf", 11)
F_BODY  = f("Arial.ttf", 12.5)
F_SMALL = f("Arial.ttf", 11.5)
F_NOTE  = f("Arial.ttf", 12)
F_HEAD  = f("Arial Bold.ttf", 11.5)
F_LOOP  = f("Arial Bold.ttf", 12)

img = Image.new("RGB", (W, H), (255, 255, 255))
d = ImageDraw.Draw(img)


def box(x, y, w, h, fill, outline, r=10, width=1.5):
    d.rounded_rectangle([x * S, y * S, (x + w) * S, (y + h) * S],
                        radius=r * S, fill=fill, outline=outline, width=int(width * S))


def txt(x, y, s, font, fill):
    d.text((x * S, y * S), s, font=font, fill=fill, anchor="ls")


def line(x1, y1, x2, y2, color, width=2, dash=None):
    pts = [(x1, y1), (x2, y2)]
    if dash is None:
        d.line([[p[0] * S, p[1] * S] for p in pts], fill=color, width=int(width * S))
        return
    dash_, gap = dash
    # walk the polyline emitting dashes
    import math
    cur = pts[0]
    for nxt in pts[1:]:
        dx, dy = nxt[0] - cur[0], nxt[1] - cur[1]
        seg = math.hypot(dx, dy)
        if seg == 0:
            continue
        ux, uy = dx / seg, dy / seg
        t = 0.0
        while t < seg:
            a = t
            b = min(t + dash_, seg)
            d.line([(cur[0] + ux * a) * S, (cur[1] + uy * a) * S,
                    (cur[0] + ux * b) * S, (cur[1] + uy * b) * S],
                   fill=color, width=int(width * S))
            t = b + gap
        cur = nxt


def head(x, y, color, direction="right", size=7):
    h = size
    if direction == "right":
        d.polygon([(x * S, y * S), ((x - h) * S, (y - h / 1.8) * S), ((x - h) * S, (y + h / 1.8) * S)], fill=color)
    elif direction == "up":
        d.polygon([(x * S, y * S), ((x - h / 1.8) * S, (y + h) * S), ((x + h / 1.8) * S, (y + h) * S)], fill=color)


def arrow(x1, y1, x2, y2, color, width=2):
    d.line([x1 * S, y1 * S, x2 * S, y2 * S], fill=color, width=int(width * S))
    if y1 == y2:
        head(x2, y2, color, "right")
    else:
        head(x1, y2, color, "up" if y2 < y1 else "right")


# ── the pass: six stages ─────────────────────────────────────────────────
stages = [
    (30,  "ASK",      ["a question,", "in plain language"], None),
    (212, "EMBED",    ["768-d point,", "on the sphere"], None),
    (394, "ROUTE",    ["argmax over", "13 anchors —"], "a label, not a weight"),
    (576, "RETRIEVE", ["HNSW · top 50 ·", "flat across cubes"], None),
    (758, "SCORE",    ["querySim,", "sorted descending"], "floor 0.55–0.65"),
    (940, "GROUND",   ["best similarity", "≥ 0.750?"], None),
]
for x, label, body, sub in stages:
    box(x, 150, 152, 104, FILL_STAGE, FRAME)
    txt(x + 16, 176, label, F_LABEL, ACCENT)
    txt(x + 16, 202, body[0], F_BODY, INK)
    txt(x + 16, 220, body[1], F_BODY, INK)
    if sub:
        txt(x + 16, 238, sub, F_SMALL, MUTED)

for x in (182, 364, 546, 728, 910):
    arrow(x, 202, x + 24, 202, ACCENT)

# ── branch ───────────────────────────────────────────────────────────────
arrow(1092, 202, 1118, 202, ACCENT)
d.line([1118 * S, 202 * S, 1118 * S, 160 * S], fill=GREEN, width=2 * S)
arrow(1118, 160, 1142, 160, GREEN)
d.line([1118 * S, 202 * S, 1118 * S, 258 * S], fill=RED, width=2 * S)
arrow(1118, 258, 1142, 258, RED)
txt(1121, 152, "yes", f("Arial Bold.ttf", 11), GREEN)
txt(1121, 277, "no", f("Arial Bold.ttf", 11), RED)

box(1146, 132, 148, 58, FILL_TERM, GREEN)
txt(1160, 155, "SPEAK", F_LABEL, GREEN)
txt(1160, 175, "the reader phrases", F_BODY, INK)

box(1146, 230, 148, 58, FILL_REFUSE, RED)
txt(1160, 253, "DECLINE", F_LABEL, RED)
txt(1160, 273, "\u201cnot in the store\u201d", F_BODY, INK)

# ── the store ────────────────────────────────────────────────────────────
box(30, 346, 1264, 150, FILL_STORE, FRAME, r=12)
txt(52, 374, "THE STORE — READ-ONLY DURING THE PASS", F_HEAD, MUTED)

txt(52, 404, "documents → chunks → entries", F_BODY, INK)
txt(52, 424, "each entry carries provenance edges back to its chunk and its source document", F_SMALL, MUTED)
txt(52, 456, "13 routing anchors = 6 antipodal axes + a residual", F_BODY, INK)
txt(52, 476, "the anchor matrix has rank 7 — the names are poles, not 13 separable modes", F_SMALL, MUTED)

txt(700, 404, "bridge edges, computed offline", F_BODY, INK)
txt(700, 424, "cross-domain links between chunks close in vector space —", F_SMALL, MUTED)
txt(700, 444, "enrichment for the reader, never a routing decision", F_SMALL, MUTED)
txt(700, 476, "vectors in LanceDB · workflow state in SQLite", F_BODY, INK)

# read-only connectors
d.line([470 * S, 344 * S, 470 * S, 262 * S], fill=GRAY, width=int(1.5 * S))
head(470, 258, GRAY, "up")
txt(482, 306, "the pass reads", F_NOTE, MUTED)

d.line([834 * S, 344 * S, 834 * S, 262 * S], fill=GRAY, width=int(1.5 * S))
head(834, 258, GRAY, "up")
txt(846, 306, "nothing writes back", F_NOTE, MUTED)

# ── no recurrence ────────────────────────────────────────────────────────
line(1220, 512, 1220, 556, PALE, 1.5, dash=(6, 5))
line(1220, 556, 100, 556, PALE, 1.5, dash=(6, 5))
line(100, 556, 100, 512, PALE, 1.5, dash=(6, 5))

d.ellipse([(660 - 15) * S, (556 - 15) * S, (660 + 15) * S, (556 + 15) * S],
          fill=(255, 255, 255), outline=PALE, width=int(1.5 * S))
d.line([652 * S, 548 * S, 668 * S, 564 * S], fill=PALE, width=int(2.2 * S))
d.line([668 * S, 548 * S, 652 * S, 564 * S], fill=PALE, width=int(2.2 * S))

d.text((660 * S, 596 * S), "NO RECURRENCE — ONE DIRECTION, NO SECOND LOOK",
       font=F_LOOP, fill=GRAY, anchor="ms")

out = "forward-pass-on-the-graph.png"
img = img.crop((0, 110 * S, W, H))          # trim the empty top margin
img.save(out, "PNG", optimize=True)
print("wrote", out, img.size)