"""Transparent chapter 00 examples. Python 3; no packages or network."""
from math import exp, log, sqrt
import json
import re
import unicodedata

VOCAB = ['<unk>', 'кошка', 'сидит', 'на', 'ковре', 'окне',
         'собака', 'спит', 'рядом', 'с', 'дом', 'улице']
E = [[0,0],[1,0],[0,1],[1,1],[1,-1],[-1,1],[-1,0],[0,-1],
     [.5,.5],[-.5,.5],[.5,-.5],[-.5,-.5]]
U = [[0,1,0,1,1,-1,-1,0,.5,-.5,.5,-.5],
     [0,0,1,1,-1,1,0,-1,.5,.5,-.5,-.5]]
BIAS = [-3,0,0,0,1,.6,0,0,0,0,0,.2]

def softmax(values):
    maximum = max(values)
    terms = [exp(v-maximum) for v in values]
    return [v/sum(terms) for v in terms]

def next_token(text='кошка', temperature=1, seed=1):
    if len(text) > 160 or not .1 <= temperature <= 2:
        raise ValueError('Input limit or temperature range')
    if not isinstance(seed, int) or not 1 <= seed <= 2147483646:
        raise ValueError('Seed must be an integer in 1..2147483646')
    # Unicode category L/N matches the browser's letter/number tokenizer.
    normalized = unicodedata.normalize('NFC', text).lower()
    normalized = ''.join(c if unicodedata.category(c)[0] in 'LN' else ' ' for c in normalized)
    groups = normalized.split()
    if not 1 <= len(groups) <= 12:
        raise ValueError('Expected 1..12 tokens')
    ids = [VOCAB.index(w) if w in VOCAB else 0 for w in groups]
    x = [[E[t][0]+i/10, E[t][1]-i/10] for i,t in enumerate(ids)]
    # Identity projections: Q=K=V=X. Causal mask is applied before softmax.
    scores = [[sum(a*b for a,b in zip(q,k))/sqrt(2) if j<=i else -float('inf')
               for j,k in enumerate(x)] for i,q in enumerate(x)]
    attention = [softmax(row) for row in scores]
    mixed = [[sum(attention[i][j]*x[j][d] for j in range(len(x)))
              for d in range(2)] for i in range(len(x))]
    residual = [[x[i][d]+mixed[i][d] for d in range(2)] for i in range(len(x))]
    # W1=I, W2=.5I, zero biases, second residual. No layer normalization.
    hidden = [[v+.5*max(0,v) for v in row] for row in residual]
    h = hidden[-1]
    logits = [h[0]*U[0][j]+h[1]*U[1][j]+BIAS[j] for j in range(12)]
    probabilities = softmax([z/temperature for z in logits])
    next_seed = (16807*seed) % 2147483647
    u = (next_seed-1)/2147483646
    cumulative = 0
    chosen = 11
    for j,p in enumerate(probabilities):
        cumulative += p
        if u < cumulative:
            chosen = j
            break
    return {'ids':ids,'A':attention,'R':residual,'H':hidden,
            'z':logits,'pi':probabilities,'u':u,'chosen':chosen}

def train_step(eta=.1):
    x,w,b,target = 2,1,0,6
    y = w*x+b
    error = y-target
    loss = .5*error**2
    gw,gb = error*x,error
    new_w,new_b = w-eta*gw,b-eta*gb
    new_y = new_w*x+new_b
    return {'y':y,'loss':loss,'gradient':[gw,gb],
            'updated':[new_w,new_b],'new_y':new_y,'new_loss':.5*(new_y-target)**2}

if __name__ == '__main__':
    print(json.dumps({'training':train_step(),
                      'exact_softmax':softmax([0,log(2),log(3)]),
                      'cat':next_token()}, ensure_ascii=False, indent=2))