"""Small explicit analytic SGD learner. No symbolic scoring/oracle dependency."""
import copy
import numpy as np
from builder import build
PORTS={'address_positions':8,'context_positions':8,'clock':'external-step'}

class Branch:
    def __init__(self,m,r,seed,addresses=None,zero=False):
        self.manifest=build(m,r,PORTS,seed);p=self.manifest['parameters'];self.W=np.array(p['W']);self.b=np.array(p['b']);self.V=np.array(p['V']);self.h=np.zeros(m);self.delay=np.zeros(m);self.addresses=np.array([i%8 for i in range(m)] if addresses is None else addresses)
        if zero:self.V[:]=0
    def reset(self):self.h[:]=0;self.delay[:]=0
    def features(self,x):return np.tanh(self.W@x+self.b)
    def query(self,x):
        if not np.array_equal(self.h,self.delay):raise ValueError('Delay/state boundary invariant violated')
        phi=self.features(x);return float(self.h@self.V@phi),phi
    def write(self,a,v):
        w=(self.addresses==a).astype(float);self.h=(1-w)*self.delay+w*np.tanh(v);self.delay=self.h.copy()
    def grads(self,x,phi,error):
        feature_error=error*(self.V.T@self.h)*(1-phi*phi)
        return {'V':error*np.outer(self.h,phi),'W':np.outer(feature_error,x),'b':feature_error}
    def cost(self,backward=False):
        m,r=self.V.shape
        # Includes multiply-adds and conservative elementwise-equivalent charges.
        forward=r*32+r*3+m*r+r+m*5
        backward_cost=m*r*2+r*32+r*8
        return forward+(backward_cost if backward else 0)

class Model:
    def __init__(self,m=8,r=8,seed=1):self.branches=[Branch(m,r,seed)];self.tick=0
    def reset(self):
        for b in self.branches:b.reset()
        self.tick=0
    def step(self,event,K,C,ablate=False):
        if not 1<=K<=8 or not 1<=C<=8:raise ValueError('Schema bounds')
        if event['type']=='WRITE':
            if not 0<=event['a']<K or event['v'] not in (-1,1):raise ValueError('Inactive/invalid write')
            for b in self.branches:b.write(event['a'],event['v'])
            result=(0.0,None,[])
        elif event['type']=='QUERY':
            if not 0<=event['c']<C or not 0<=event['j']<K:raise ValueError('Inactive query')
            x=np.zeros(32);x[event['c']]=1;x[8+event['j']]=1;x[16:16+K]=1;x[24:24+C]=1
            features=[b.query(x) for b in self.branches];result=(sum(y for i,(y,f) in enumerate(features) if not(ablate and i>0)),x,[f for y,f in features])
        else:raise ValueError('Unknown event')
        self.tick+=1;return result
    def snapshot(self):return copy.deepcopy(self)
    def grow(self,seed):
        child=self.snapshot();child.branches.append(Branch(4,4,seed,addresses=[0,1,2,3],zero=True));return child
    def parameters(self):return sum(b.W.size+b.b.size+b.V.size for b in self.branches)

def run_episode(model,events,answers,K,C,train=False,lr=.08,ablate=False):
    model.reset();gradient=[{k:np.zeros_like(getattr(b,k)) for k in ['W','b','V']} for b in model.branches];loss=0.;correct=0;queries=0;mac=0
    for e,target in zip(events,answers):
        prediction,x,features=model.step(e,K,C,ablate);mac+=sum(b.cost(train and e['type']=='QUERY') for b in model.branches)
        if e['type']=='QUERY':
            error=prediction-target;loss+=error*error/2;queries+=1;correct+=int(np.sign(prediction)==target)
            if train:
                for b,phi,g in zip(model.branches,features,gradient):
                    for k,v in b.grads(x,phi,error).items():g[k]+=v
    if train:
        for b,g in zip(model.branches,gradient):
            for k,v in g.items():setattr(b,k,getattr(b,k)-lr*v/queries)
    if not all(np.isfinite(getattr(b,k)).all() for b in model.branches for k in ['W','b','V']):raise ValueError('Nonfinite learner')
    return {'accuracy':correct/queries,'loss':loss/queries,'queries':queries,'mac_equivalents':mac,'events':len(events)}