"""Streaming inference port of the supplied causal byte NCA checkpoint.
Dense matrix operations: zipp_gpu (WebGL2 or explicit CPU reference).
Nonlinearities and cache bookkeeping: Python in ZIPP WASM.
No attention, no fast-memory branch, no online parameter training.
"""
import math
import _zipp_gpu
W = {}
S = {}
cache = []
state = []
anchor = []
new_cache = []
stage = 0

def configure(model):
    global W
    W = model['parameters']
    reset(1)
    return True

def reset(seed):
    global S, cache
    cache = [[] for _ in range(8)]
    S = {'mode':'language','bytes':[],'logits':[0.0]*256,'stages':[],
         'busy':False,'error':None,'graphs':0,'lastMs':0,'seed':max(1,int(seed)%2147483647),
         'message':'Start with a prompt. The cache is empty; trained weights are frozen.'}

def failed(error):
    S['error'] = str(error); S['busy'] = False

def mark(result):
    S['graphs'] += 1; S['lastMs'] = result['stats']['totalWallMs']

def submit_dense(layers, callback):
    # The immutable checkpoint arrays already contain finite float32 values.
    # Build the same data-only graph protocol directly, avoiding repeated
    # Python struct.pack/unpack of thousands of unchanged weights. ZIPP's host
    # validates every input, operation, bound and output before execution.
    nodes = []
    outputs = []
    for name, weight, bias, data, rows, cols in layers:
        base = len(nodes)
        nodes.append({'id':base,'op':'input','shape':[rows,cols],'data':weight})
        nodes.append({'id':base+1,'op':'input','shape':[cols,1],'data':data})
        nodes.append({'id':base+2,'op':'matmul','a':base,'b':base+1})
        nodes.append({'id':base+3,'op':'input','shape':[rows,1],'data':bias})
        nodes.append({'id':base+4,'op':'add','a':base+2,'b':base+3})
        outputs.append({'name':name,'id':base+4})
    def delivered(reply):
        if reply.get('ok'):
            callback(reply['value'])
        else:
            failed(reply.get('error',{}).get('message','Compute graph failed'))
    _zipp_gpu.post({'version':1,'nodes':nodes,'outputs':outputs},delivered)

def layer(name, data, rows, cols):
    return (name,W[name+'.weight']['data'],W[name+'.bias']['data'],data,rows,cols)

def feed(token):
    global state, anchor, new_cache, stage
    if S['busy']: raise ValueError('A byte is already being processed')
    if not 0 <= token < 256: raise ValueError('Byte outside range')
    S['busy'] = True
    S['bytes'].append(token)
    state = [math.tanh(x) for x in W['embedding.weight']['data'][token*32:(token+1)*32]]
    new_cache = []; stage = 0; S['stages'] = []
    def done(result):
        global anchor
        mark(result)
        anchor = result['outputs']['anchor']['data']
        update_stage()
    submit_dense([layer('anchor',state,64,32)],done)

def update_stage():
    global new_cache
    if stage == 8:
        finish_byte()
        return
    # Each stage stores predecessor states BEFORE that stage's update.
    columns = cache[stage] + [state[:]]
    new_cache.append(columns[-2:])
    window = [[0.0]*32 for _ in range(3-len(columns))] + columns
    flat = [window[j][i] for i in range(32) for j in range(3)]
    def perceived(result):
        mark(result)
        values = result['outputs']['h']['data']
        hvalues = [x/(1.0+math.exp(max(-60.0,min(60.0,-x)))) for x in values]
        submit_dense([layer('gate',hvalues,32,64),layer('candidate',hvalues,32,64)],updated)
    submit_dense([('h',W['perception.weight']['data'],anchor,flat,64,96)],perceived)

def updated(result):
    global state, stage
    mark(result)
    gates = result['outputs']['gate']['data']
    candidates = result['outputs']['candidate']['data']
    next_state = []
    for i in range(32):
        gate = 1.0/(1.0+math.exp(max(-60.0,min(60.0,-gates[i]))))
        next_state.append((1.0-gate)*state[i]+gate*math.tanh(candidates[i]))
    state = next_state
    S['stages'].append(state[:]); stage += 1
    update_stage()

def finish_byte():
    global cache
    def done(result):
        global cache
        mark(result)
        S['logits'] = result['outputs']['head']['data']
        cache = new_cache
        S['busy'] = False
        S['message'] = 'Next-byte logits computed using eight causal updates.'
    submit_dense([layer('head',state,256,32)],done)

def act(action,args):
    if S['busy']: raise ValueError('Wait for the current byte')
    S['error'] = None
    if action == 'reset':
        reset(args.get('seed',1))
    elif action == 'feed':
        if len(S['bytes']) >= 384: raise ValueError('Session limit reached; reset the model')
        feed(int(args['token']))
    elif action == 'sample':
        if len(S['bytes']) == 0: raise ValueError('Feed a prompt first')
        if len(S['bytes']) >= 384: raise ValueError('Session limit reached; reset the model')
        temperature = float(args.get('temperature',0.3))
        if not 0.1 <= temperature <= 2.0: raise ValueError('Temperature must be 0.1–2.0')
        logits = S['logits']; peak = max(logits)
        probabilities = [math.exp((x-peak)/temperature) for x in logits]
        S['seed'] = (S['seed']*48271)%2147483647
        target = (S['seed']/2147483647.0)*sum(probabilities)
        token = 255; acc = 0.0
        for i in range(256):
            acc += probabilities[i]
            if acc >= target:
                token = i; break
        feed(token)
    else:
        raise ValueError('Unknown language command')
    return snapshot()

def snapshot():
    return S