"""Generate original static SVGs from the same inputs as the Article 06 videos.

Run this file from any directory. --check compares without modifying the SVGs.
Only standard Python and NumPy are required; no trained-model data are used.
"""
import argparse
from html import escape
from pathlib import Path
import xml.etree.ElementTree as ET
from model import checkpoint

ROOT = Path(__file__).resolve().parents[3]
OUT = ROOT / 'assets/img/posts/first-principles-06'
D = checkpoint()
INK, MUTED, BG = '#15243b', '#475569', '#f8fafc'
Q, K, V, P = '#1d4ed8', '#047857', '#b45309', '#6d28d9'


def text(x, y, s, size=30, color=INK, anchor='start', weight=400):
    return f'<text x="{x}" y="{y}" font-size="{size}" fill="{color}" text-anchor="{anchor}" font-weight="{weight}">{escape(str(s))}</text>'


def rect(x, y, w, h, color='#cbd5e1', fill='white', radius=16):
    return f'<rect x="{x}" y="{y}" width="{w}" height="{h}" rx="{radius}" fill="{fill}" stroke="{color}" stroke-width="2"/>'


def line(x1, y1, x2, y2, color=MUTED, arrow=True, width=3):
    marker = f' marker-end="url(#arrow-{color[1:]})"' if arrow else ''
    return f'<path d="M {x1} {y1} L {x2} {y2}" fill="none" stroke="{color}" stroke-width="{width}"{marker}/>'


def svg(title, desc, body):
    markers = ''.join(f'<marker id="arrow-{c[1:]}" viewBox="0 0 10 10" refX="9" refY="5" markerWidth="4" markerHeight="4" orient="auto-start-reverse"><path d="M 0 0 L 10 5 L 0 10 z" fill="{c}"/></marker>' for c in (INK, MUTED, Q, K, V, P, '#cbd5e1'))
    s = f'''<svg xmlns="http://www.w3.org/2000/svg" width="1200" height="675" viewBox="0 0 1200 675" role="img" aria-labelledby="title desc">
<title id="title">{escape(title)}</title><desc id="desc">{escape(desc)}</desc>
<defs>{markers}</defs>
<rect width="1200" height="675" rx="20" fill="{BG}"/>
<g font-family="Arial, Helvetica, sans-serif">{text(44,62,title,38,weight=700)}{body}</g></svg>\n'''
    ET.fromstring(s)
    return s


def matrix(cx, y, a, color, name, fmt='.0f', highlight=None):
    result = text(cx, y-62, name, 32, color, 'middle', 700)
    rows, cols = a.shape
    cw, rh = 86, 58
    left = cx-cols*cw/2
    for i in range(rows):
        if i == highlight:
            result += rect(left-8,y+i*rh-37,cols*cw+16,rh,P,'#f3e8ff',8)
        for j in range(cols):
            result += text(left+(j+.5)*cw,y+i*rh,format(float(a[i,j]),fmt),29,color,'middle')
    result += line(left-16,y-42,left-16,y+(rows-1)*rh+16,color,False,3)
    result += line(left+cols*cw+16,y-42,left+cols*cw+16,y+(rows-1)*rh+16,color,False,3)
    return result


def build():
    figures = {}
    body = text(44,108,'Illustrative records about Felix’s notebook; not measured model behavior.',25,MUTED)
    body += text(180,190,'QUERY',27,Q,'middle',700)+text(580,190,'KEY',27,K,'middle',700)+text(995,190,'VALUE',27,V,'middle',700)
    body += rect(44,260,280,150,Q,'#eff6ff')+text(184,316,'Where is',34,Q,'middle',700)+text(184,365,'the notebook?',34,Q,'middle',700)
    for y, clue, content in [(235,'notebook · location','meeting room'),(400,'notebook · color','blue')]:
        body += rect(400,y,750,132)
        body += text(580,y+74,clue,30,K,'middle',600)+text(995,y+74,content,32,V,'middle',600)
        body += line(791,y+24,791,y+108,'#cbd5e1',False,2)
        body += line(324,335,390,y+66,Q)
    body += text(44,595,'Match using the Key; retrieve the paired Value.',30,weight=700)
    body += text(44,637,'The Key does not generate the Value. Both can be derived from the same input.',25,MUTED)
    figures['01-qkv-roles.svg'] = svg('Finding a record is not retrieving its content', 'A request about notebook location is compared with matching clues. Each record pairs a Key with a Value. The labels are an analogy, not learned semantic dimensions.',body)

    body = text(44,111,'Column-vector view: one input, three shared linear maps.',27,MUTED)
    body += rect(44,200,240,100)+text(164,262,'x₃ = (1, 1)ᵀ',34,INK,'middle',600)
    for y,c,w,result in [(177,Q,'W_Qᵀ x₃','q₃ = (1, 1)ᵀ'),(287,K,'W_Kᵀ x₃','k₃ = (1, 2)ᵀ'),(397,V,'W_Vᵀ x₃','v₃ = (2, 1)ᵀ')]:
        body += line(290,250,414,y,c)+rect(425,y-38,300,76,c)+text(575,y+10,w,34,c,'middle',600)
        body += line(735,y,825,y,c)+text(1000,y+10,result,34,c,'middle',600)
    body += rect(44,495,1110,122)
    body += text(600,539,'Stack the position vectors as rows of X:',28,INK,'middle')
    body += text(240,588,'Q = XW_Q',32,Q,'middle',600)+text(600,588,'K = XW_K',32,K,'middle',600)+text(960,588,'V = XW_V',32,V,'middle',600)
    body += text(44,653,'Each position has its own vectors; the projection parameters are shared.',25,MUTED)
    figures['02-qkv-projections.svg'] = svg('One position, then the whole sequence','For the third input (1,1), the article matrices produce q3=(1,1), k3=(1,2), and v3=(2,1). Batched row conventions are Q=XWQ, K=XWK and V=XWV.',body)

    body = text(44,111,'Rows: query positions. Columns: key positions. Query 2 is highlighted.',27,MUTED)
    body += matrix(210,250,D['raw'],INK,'QKᵀ',highlight=1)
    body += matrix(600,250,D['scores'],Q,'E = QKᵀ / √2','.3f',1)
    body += matrix(990,250,D['weights'],K,'A = Softmax_rows(E)','.3f',1)
    body += line(380,302,420,302,Q)+line(770,302,810,302,K)
    body += text(600,472,'Second row: [0.248255, 0.248255, 0.503490]',32,P,'middle',700)
    body += text(600,535,'Every row of A sums to 1.',32,INK,'middle')
    body += text(600,591,'These weights mix source Values; they are not vocabulary probabilities.',25,MUTED,'middle')
    figures['03-attention-scores-to-weights.svg'] = svg('Scores become normalized reading weights','All nine entries are computed from the article matrices. Raw scores are divided by square root of two in every row, then Softmax is applied across each row.',body)

    def pt(a): return (120+180*float(a[0]),510-180*float(a[1]))
    body = text(44,111,'Same Values, fixed parameters; query 2 supplies the mixing weights.',27,MUTED)
    ox,oy=pt((0,0))
    body += line(ox,oy,565,oy,arrow=True)+line(ox,oy,ox,230,arrow=True)
    body += text(330,572,'coordinate 1',24,MUTED,'middle')+text(60,218,'coordinate 2',24,MUTED)
    body += '<path d="M 480 510 L 120 330 L 480 330 Z" fill="#e2e8f0" fill-opacity=".55" stroke="#94a3b8" stroke-dasharray="7 6"/>'
    for a,name,dx,dy in [((2,0),'v₁=(2,0)',16,28),((0,1),'v₂=(0,1)',10,-18),((2,1),'v₃=(2,1)',0,-18)]:
        xx,yy=pt(a); body += f'<circle cx="{xx}" cy="{yy}" r="6" fill="{MUTED}"/>'+text(xx+dx,yy+dy,name,24,MUTED)
    cur=D['v'][0]*0
    for j,c in enumerate((Q,K,V)):
        end=cur+D['weights'][1,j]*D['v'][j]
        body += line(*pt(cur),*pt(end),color=c,width=5)
        cur=end
    xx,yy=pt(D['out'][1]); body += f'<circle cx="{xx}" cy="{yy}" r="9" fill="{P}"/>'+text(xx-8,yy-20,'o₂',31,P,'end',700)
    for y,c,s in [(245,Q,'0.248255 v₁ = (0.496510, 0)'),(315,K,'0.248255 v₂ = (0, 0.248255)'),(385,V,'0.503490 v₃ = (1.006980, 0.503490)')]:
        body+=text(650,y,s,25,c,weight=600)
    body+=rect(635,440,520,100,P,'#f3e8ff')+text(895,501,'o₂ ≈ (1.503490, 0.751745)',28,P,'middle',700)
    body+=text(44,611,'The colored arrows add weighted Values head-to-tail.',29,weight=600)
    body+=text(44,652,'The output lies in the convex hull of these Values. It is not a word prediction.',25,MUTED)
    figures['04-weighted-value-sum.svg'] = svg('The weights say how much; Values say what','The three weighted vectors are added head-to-tail. Their sum is (1.503490,0.751745), inside the triangle formed by the three Values. This is a constructed geometric example.',body)

    body=text(44,111,'Compare all-visible and causal reads using the same Q, K, and V.',27,MUTED)
    body+=rect(44,165,540,370,Q)+rect(616,165,540,370,P)
    for x,title,c in [(314,'ALL VISIBLE',Q),(886,'CAUSAL: ROW 2',P)]:
        body+=text(x,214,title,30,c,'middle',700)
    for y,left,right in [(278,'scores: [0.707, 0.707, 1.414]','scores: [0.707, 0.707, −∞]'),(357,'weights: [0.248, 0.248, 0.503]','weights: [0.5, 0.5, 0]'),(455,'o₂ ≈ (1.503490, 0.751745)','o₂ = (1, 0.5)')]:
        body+=text(314,y,left,27,Q,'middle',600)+text(886,y,right,27,P,'middle',600)
    body+=text(314,400,'sum = 1',25,MUTED,'middle')+text(886,400,'sum = 1',25,MUTED,'middle')
    body+=text(44,585,'Mask before Softmax: the visible candidates receive all the weight.',29,weight=700)
    body+=text(44,634,'Position i has received tᵢ; its output helps predict tᵢ₊₁. It may read j ≤ i.',26,MUTED)
    figures['05-causal-mask.svg'] = svg('Masking changes the normalization set','For query 2, future key 3 is assigned minus infinity before Softmax. The two remaining equal scores receive weights one half each. The output changes from (1.503490,0.751745) to (1,0.5).',body)
    return figures


def main():
    parser=argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--check',action='store_true')
    args=parser.parse_args()
    OUT.mkdir(parents=True,exist_ok=True)
    for name,content in build().items():
        path=OUT/name
        if args.check:
            assert path.read_text(encoding='utf-8') == content, f'Outdated figure: {name}'
        else:
            path.write_text(content,encoding='utf-8')
    print('PASS: five reproducible, XML-parseable Article 06 figures')


if __name__=='__main__':
    main()