"""Browser inference port of CellularMemory (fast variant, supplied seed-0 weights).
Python runs in ZIPP WASM; all write/read/decoder matrix arithmetic is submitted
through zipp_gpu. Observed labels are used ONLY by the interface, never queries.
This is not a re-training experiment or a claim of learned routing.
"""
import math
from zipp_gpu import Graph

W = {}
S = {}
queue = []
keys = []

def configure(model):
    global W, keys
    W = model['parameters']
    raw = W['keys.weight']['data']
    keys = []
    for i in range(16):
        k = raw[i*16:(i+1)*16]
        norm = max(math.sqrt(sum(x*x for x in k)), 1e-8)
        keys.append([x/norm for x in k])
    clear()
    return True

def clear():
    global S
    S = {'mode':'memory','matrix':[[0.0]*128 for _ in range(8)],
         'records':[], 'key':5,'origin':6,'position':6,'steps':0,
         'acc':[0.0]*8,'logits':[0.0]*8,'answer':None,'relay':True,
         'querying':False,'busy':False,'error':None,'graphs':0,'lastMs':0,
         'message':'Empty private memory. Shared checkpoint weights are frozen.'}

def failed(error):
    S['error'] = str(error)
    S['busy'] = False
    queue.clear()

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

def _write_next():
    if len(queue) == 0:
        S['busy'] = False
        S['message'] = 'Observed associations written locally. Shared weights unchanged.'
        return
    event = queue.pop(0)
    k = keys[event['key']]
    cell = event['cell']
    g = Graph()
    a = g.tensor(S['matrix'][cell], shape=[8,16])
    kc = g.tensor(k, shape=[16,1])
    value = g.tensor([float(b*2-1) for b in event['bits']], shape=[8,1])
    enc = g.tensor(W['value_encoder.weight']['data'], shape=[8,8])
    error = enc @ value - a @ kc
    new = a + (error @ kc.transpose()) * (1.0 / (1e-8 + sum(x*x for x in k)))
    def complete(result):
        mark(result)
        S['matrix'][cell] = result['outputs']['memory']['data']
        S['records'] = [r for r in S['records'] if r['key'] != event['key']]
        S['records'].append(event)
        _write_next()
    g.submit(complete, failed, memory=new)

def act(action, args):
    global queue
    if S['busy']:
        raise ValueError('Wait for the current graph to finish')
    S['error'] = None
    if action == 'reset':
        clear()
        seed = int(args.get('seed', 1)) % 2147483647
        events = []
        for i in range(8):
            seed = (seed*48271+1) % 2147483647
            number = seed % 256
            events.append({'key':(i*5+5)%16,'cell':(i*3+2)%8,
                           'bits':[(number >> (7-j)) & 1 for j in range(8)]})
        queue = events
        S['busy'] = True
        _write_next()
    elif action == 'clear':
        clear()
    elif action == 'write':
        key = int(args['key']); cell = int(args['cell']); bits = args['bits']
        if not 0 <= key < 16 or not 0 <= cell < 8 or len(bits) != 8 or any(b not in [0,1] for b in bits):
            raise ValueError('Use a key 0–15, cell 0–7 and eight binary digits')
        for record in S['records']:
            if record['key'] == key and record['cell'] != cell:
                raise ValueError('Overwrite this key at its original writer cell, or reset first')
        S['querying'] = False; S['answer'] = None
        queue = [{'key':key,'cell':cell,'bits':bits}]
        S['busy'] = True
        _write_next()
    elif action == 'query':
        key = int(args.get('key',5)); origin = int(args.get('origin',6))
        if not 0 <= key < 16 or not 0 <= origin < 8:
            raise ValueError('Query key/cell outside allowed range')
        S['key'] = key; S['origin'] = origin; S['position'] = origin
        S['steps'] = 0; S['acc'] = [0.0]*8; S['logits'] = [0.0]*8
        S['answer'] = None; S['querying'] = True
        S['message'] = 'Fresh query packet. No observed value is included in the query.'
    elif action == 'relay':
        S['relay'] = bool(args['enabled'])
    elif action == 'erase':
        cell = int(args.get('cell',-1))
        if cell < -1 or cell >= 8: raise ValueError('Invalid erase cell')
        if cell == -1:
            S['matrix'] = [[0.0]*128 for _ in range(8)]
        else:
            S['matrix'][cell] = [0.0]*128
        S['querying'] = False; S['answer'] = None
        S['message'] = 'Private matrix erased. Labels remain as a reference, not model input.'
    elif action == 'step':
        if not S['querying'] or S['steps'] >= 8:
            return snapshot()
        S['busy'] = True
        if S['relay']:
            S['position'] = (S['position']+1)%8
        g = Graph()
        a = g.tensor(S['matrix'][S['position']],shape=[8,16])
        k = g.tensor(keys[S['key']],shape=[16,1])
        acc = g.tensor(S['acc'],shape=[8,1]) + a @ k
        decoder = g.tensor(W['decoder.weight']['data'],shape=[8,8])
        logits = (decoder @ acc) * math.exp(max(-2.0,min(5.0,W['log_scale']['data'][0])))
        def complete(result):
            mark(result)
            S['acc'] = result['outputs']['acc']['data']
            S['steps'] += 1
            if S['steps'] == 8:
                S['logits'] = result['outputs']['logits']['data']
                S['answer'] = [1 if x > 0 else 0 for x in S['logits']]
                S['querying'] = False
                S['message'] = 'Query complete. The answer was decoded from private memory.'
            S['busy'] = False
        g.submit(complete,failed,acc=acc,logits=logits)
    else:
        raise ValueError('Unknown memory command')
    return snapshot()

def snapshot():
    return S