"""Own offline MLP; export integer AltAILIF hidden rates and native two-neuron readout."""
from pathlib import Path
import numpy as np,json,hashlib,time
R=Path.home()/'run/pets-knp';M=R/'head-next';M.mkdir(exist_ok=False);rng=np.random.default_rng(20261015)
a=np.load(R/'training-counts.npz');yt=np.tile(np.load(R/'data/train.npz')['labels'],2).astype(np.float32);yv=np.load(R/'data/val.npz')['labels']
def encode(x):
 q=np.sqrt(x.astype(np.float64));q*=2048/np.maximum(np.sqrt((q*q).sum(1,keepdims=True)),1)
 return np.clip(np.rint(q),0,640).astype(np.int32)
xi=encode(a['train']);vi=encode(a['validation']);mu=xi.mean(0).astype(np.float32);sd=(xi.std(0)+2).astype(np.float32);x=((xi-mu)/sd).astype(np.float32);v=((vi-mu)/sd).astype(np.float32)
h=256;params=[rng.normal(0,np.sqrt(2/800),(800,h)).astype(np.float32),np.zeros(h,np.float32),rng.normal(0,.05,h).astype(np.float32),np.array(0,np.float32)];mm=[np.zeros_like(p) for p in params];vv=[np.zeros_like(p) for p in params];step=0;history=[];best=None
for epoch in range(1,41):
 lr=.0015*(.98**(epoch-1));order=rng.permutation(len(x))
 for start in range(0,len(x),256):
  ix=order[start:start+256];xx=x[ix];yy=yt[ix];w,b,u,c=params;z=xx@w+b;act=np.clip(z,0,16);mask=(rng.random(act.shape)>.15).astype(np.float32)/.85;hidden=act*mask;scores=hidden@u+c;prob=1/(1+np.exp(-np.clip(scores,-30,30)));dz=(prob-yy)/len(xx)
  dh=dz[:,None]*u[None,:]*mask*((z>0)&(z<16));grads=[xx.T@dh+.001*w,dh.sum(0),hidden.T@dz+.001*u,np.array(dz.sum(),np.float32)];step+=1
  for k,g in enumerate(grads): mm[k]=.9*mm[k]+.1*g
  # Array/scalar Adam, avoid assigning a non-sliceable scalar.
  for k,(p,g) in enumerate(zip(params,grads)):
   vv[k]=.999*vv[k]+.001*g*g
   p-=lr*(mm[k]/(1-.9**step))/(np.sqrt(vv[k]/(1-.999**step))+1e-8)
 w,b,u,c=params;pred=(np.clip(v@w+b,0,16)@u+c>0);floatacc=float((pred==yv).mean())
 # Fold training standardization into integer input weights; per-neuron scale.
 wf=w/sd[:,None];bf=b-(mu/sd)@w
 # Bound scale BEFORE integer casts, including nearly-zero/dead units.
 assert np.isfinite(wf).all() and np.isfinite(bf).all()
 scale=np.minimum(127/np.maximum(abs(wf).max(0),1e-12),262140.)
 scale=np.minimum(scale,(2**22)/(2063*np.linalg.norm(wf.astype(np.float64),axis=0)+abs(bf)+1))
 wq=np.rint(wf*scale).astype(np.int32);bq=np.rint(bf*scale).astype(np.int32);thr=np.maximum(np.rint(scale/4),1).astype(np.int32)
 bounds=np.ceil(2063*np.linalg.norm(wq.astype(np.float64),axis=0))+abs(bq.astype(np.int64))+1
 assert thr.max()<=65535 and (2*bounds<2**24).all()

 hidden=np.clip(np.floor(np.maximum(vi.astype(np.int64)@wq+bq,0)/thr),0,64).astype(np.int32)
 # Readout refit on train hidden spike counts; offline only. Tune regularisation on validation.
 ht=np.clip(np.floor(np.maximum(xi.astype(np.int64)@wq+bq,0)/thr),0,64).astype(np.int32);hm=ht.mean(0);hs=ht.std(0)+.1;zz=(ht-hm)/hs;A=zz.T@zz;rhs=zz.T@(yt*2-1)
 epochbest=None
 for lam in (10,100,1000):
  sol=np.linalg.solve(A+lam*np.eye(h),rhs);ww=sol/hs;bb=(yt*2-1).mean()-hm@ww;ss=127/abs(ww).max();uw=np.rint(ww*ss).astype(np.int32);ub=int(round(bb*ss));margin=hidden.astype(np.int64)@uw+ub;acc=float(((margin>0)==yv).mean())
  if epochbest is None or acc>epochbest[0]:epochbest=(acc,uw,ub,lam)
 row={'epoch':epoch,'float_validation':floatacc,'native_quantized_validation':epochbest[0],'lambda':epochbest[3]};history.append(row);print(row,flush=True)
 if best is None or epochbest[0]>best:
  best=epochbest[0];np.savez(M/'float-checkpoint.npz',w=w,b=b,u=u,c=c,mean=mu,std=sd);np.savez(M/'model.npz',w=wq,b=bq,thresholds=thr,readout_w=epochbest[1],readout_b=np.array([epochbest[2]],np.int32));(M/'selection.json').write_text(json.dumps(dict(row,encoding='root_l2',hidden=h,max_hidden_count=64),indent=2))
 (M/'progress.json').write_text(json.dumps({'current':epoch,'total':40,'done':epoch==40}))
 (M/'history.json').write_text(json.dumps(history,indent=2))
# Source artifact export; no test read by this script.
m=np.load(M/'model.npz');
for key,fn in [('w','hidden_w.bin'),('b','hidden_b.bin'),('thresholds','hidden_thresholds.bin'),('readout_w','weights.bin'),('readout_b','bias.bin')]:m[key].astype(np.int32).tofile(M/fn)
vi.tofile(R/'validation-encoded.bin');print('BEST',best,flush=True)