import json, random
import numpy as np
SEED=609
np.random.seed(SEED); random.seed(SEED)

def choose_signs(cands,b,g,rho=0.05):
    n=cands.shape[0]; signs=np.ones(n); target=rho*float(g@g)
    def obj(s): return float(np.max(np.abs(b+s@cands)))
    def ok(s): return float(g@(s@cands))>=target-1e-12
    if n<=8:
        best=None
        for mask in range(1<<n):
            s=np.array([1. if (mask>>j)&1 else -1. for j in range(n)])
            if ok(s):
                v=obj(s)
                if best is None or v<best[0]: best=(v,s)
        if best is not None: return best[1]
    if not ok(signs): return signs
    cur=obj(signs)
    for _ in range(3):
        changed=False
        for j in range(n):
            q=signs.copy(); q[j]*=-1; v=obj(q)
            if ok(q) and v<cur-1e-12: signs,cur,changed=q,v,True
        if not changed: break
    return signs

def toy_scaling():
    ms=[4,8,16,32,64]; T=150; trials=8; rows=[]
    for m in ms:
        vals={k:[] for k in ('guided','all_plus','random')}
        for tr in range(trials):
            bs={k:np.zeros(m) for k in vals}
            for t in range(T):
                a=np.random.RandomState(SEED+tr*1000+t+m).randn(m-1,m)
                a/=np.maximum(np.linalg.norm(a,axis=1,keepdims=True),1e-12); g=a.mean(0)
                bs['guided']+=choose_signs(a,bs['guided'],g)@a
                bs['all_plus']+=a.sum(0); bs['random']+=np.random.choice([-1.,1.],m-1)@a
            for k in vals: vals[k].append(np.max(np.abs(bs[k])))
        rows.append({'m':m,**{k:float(np.median(v)) for k,v in vals.items()}})
    m=32; checkpoints=np.array([15,30,60,120,150]); paths={k:[] for k in ('guided','all_plus','random')}
    for tr in range(10):
        bs={k:np.zeros(m) for k in paths}
        for t in range(1,T+1):
            a=np.random.RandomState(SEED+5000+tr*1000+t).randn(m-1,m); a/=np.linalg.norm(a,axis=1,keepdims=True); g=a.mean(0)
            bs['guided']+=choose_signs(a,bs['guided'],g)@a; bs['all_plus']+=a.sum(0); bs['random']+=np.random.choice([-1.,1.],m-1)@a
            if t in checkpoints:
                for k in paths: paths[k].append((t,np.max(np.abs(bs[k]))))
    exponents={k:float(np.polyfit(np.log([z[0] for z in paths[k]]),np.log([z[1] for z in paths[k]]),1)[0]) for k in paths}
    return {'rows':rows,'growth_exponents_m32':exponents}

def neural_run(guided,steps=80,K=8,n=4):
    rng=np.random.RandomState(SEED+7); N=1024; X=rng.randn(N,2).astype('float32'); y=(X[:,0]*X[:,1]>0).astype('int64')
    p=[rng.randn(2,16).astype('float32')*.4,np.zeros(16,'float32'),rng.randn(16,2).astype('float32')*.2,np.zeros(2,'float32')]
    def grad(x,yy):
        w,bb,v,cc=p; h=np.maximum(0,x@w+bb); z=h@v+cc; z-=z.max(1,keepdims=True); pr=np.exp(z); pr/=pr.sum(1,keepdims=True); dz=(pr-np.eye(2,dtype='float32')[yy])/len(x)
        gv=h.T@dz; gc=dz.sum(0); dh=(dz@v.T)*(h>0); return np.r_[(x.T@dh).ravel(),dh.sum(0),gv.ravel(),gc]
    b=np.zeros(sum(q.size for q in p)); peaks=[]; clips=0
    for t in range(steps):
        ids=rng.choice(N,n*32,False); cs=np.stack([grad(X[ids[j*32:(j+1)*32]],y[ids[j*32:(j+1)*32]]) for j in range(n)]); g=cs.mean(0)
        s=choose_signs(cs,b,g) if guided else np.ones(n); b+=s@cs; peaks.append(np.max(np.abs(b)))
        if (t+1)%K==0:
            u=b/K
            if np.linalg.norm(u)>2: clips+=1; u*=2/np.linalg.norm(u)
            off=0
            for q in p: q-=.25*u[off:off+q.size].reshape(q.shape); off+=q.size
            b*=0
    z=np.maximum(0,X@p[0]+p[1])@p[2]+p[3]; loss=float(np.mean(np.logaddexp(0,z.max(1))-z[np.arange(N),y])); acc=float(np.mean(z.argmax(1)==y))
    return {'peak_residual_median':float(np.median(peaks)),'peak_residual_max':float(max(peaks)),'clip_events':clips,'final_accuracy':acc,'final_loss':loss}

def main():
    out={'seed':SEED,'toy':toy_scaling(),'neural':{'ordinary':neural_run(False),'guided':neural_run(True)}}
    open('results.json','w').write(json.dumps(out,indent=2)); print(json.dumps(out,indent=2))
if __name__=='__main__':main()