"""A reproducible one-head attention teaching example; Python standard library.

Inputs x1=[1,0], x2=[0,1]. Only last-position loss is illustrated.
Q/K/V each have width one; therefore sqrt(d_k)=1.
All 7 attention weights are updated by SGD, including the zero-gradient WQ[0].
The fixed two-candidate readout has logits [attention_output, 0].
This is an attention-only teaching model, not a complete Transformer.
"""
import math
import json
from pathlib import Path

def sigmoid(x):
    return 1 / (1 + math.exp(-x))

def forward(w):
    q = w['WQ1']  # x2 @ WQ
    a = sigmoid(q * (w['WK0'] - w['WK1']))
    mixed = a*w['WV0'] + (1-a)*w['WV1']
    logit = w['WO']*mixed
    p = sigmoid(logit)
    return dict(q=q, attention=[a,1-a], mixed=mixed, logit=logit,
                probability=p, loss=-math.log(p))

def backward(w):
    r=forward(w); a=r['attention'][0]
    dlogit=r['probability']-1
    dmixed=dlogit*w['WO']
    da=dmixed*(w['WV0']-w['WV1'])
    dscore=da*a*(1-a)
    return dict(WQ0=0.,WQ1=dscore*(w['WK0']-w['WK1']),
                WK0=dscore*r['q'],WK1=-dscore*r['q'],
                WV0=dmixed*a,WV1=dmixed*(1-a),WO=dlogit*r['mixed'])

def run():
    w=dict(WQ0=.2,WQ1=.5,WK0=1.,WK1=0.,WV0=1.,WV1=-1.,WO=1.)
    before=forward(w); grad=backward(w); eps=1e-6; checks={}
    for key in w:
        lo=dict(w);hi=dict(w);lo[key]-=eps;hi[key]+=eps
        numeric=(forward(hi)['loss']-forward(lo)['loss'])/(2*eps)
        checks[key]=dict(analytic=grad[key],finite_difference=numeric,
                         absolute_error=abs(numeric-grad[key]))
    lr=.1; updated={k:v-lr*grad[k] for k,v in w.items()}
    after=forward(updated)
    perturb=dict(w);perturb['WQ1']+=.001
    out=dict(scope='One head, two positions, scalar Q/K/V, last-position loss only; fixed readout.',
             inputs=[[1,0],[0,1]],weights_before=w,learning_rate=lr,
             before=before,gradients=grad,weights_after=updated,after=after,
             finite_difference_checks=checks,
             WQ1_plus_001_only=forward(perturb),
             maximum_gradient_error=max(c['absolute_error'] for c in checks.values()))
    assert out['maximum_gradient_error'] < 1e-8
    assert after['loss']<before['loss']
    Path(__file__).with_name('attention-learning-numerics.json').write_text(json.dumps(out,indent=2)+'\n')
    print(json.dumps(out,indent=2))

if __name__=='__main__':run()