"""Train integer KNP readout on train only; choose lambda/quantization on validation."""
from pathlib import Path
import numpy as np,json,re,hashlib
R=Path.home()/'run/pets-knp'; C=R/'counts';M=R/'readout-v2'
if M.exists():raise RuntimeError('Refuse to overwrite readout')
M.mkdir()
def load(tag,n):
 parts=[];nextrow=0
 for p in sorted(C.glob(tag+'_*_counts.bin'),key=lambda p:int(p.stem.split('_')[-2])):
  if not re.fullmatch(re.escape(tag)+r'_\d+_counts.bin',p.name):continue
  first=int(p.stem.split('_')[-2]);assert first==nextrow and (C/f'{tag}_{first}.done').exists()
  a=np.fromfile(p,np.int32).reshape(-1,800);parts.append(a);nextrow+=len(a)
 assert nextrow==n,(tag,nextrow,n)
 return np.concatenate(parts)
xt=np.concatenate([load('train',8000),load('train_flip',8000)]);xv=load('val',2000)
y=np.load(R/'data/train.npz')['labels'];yt=np.tile(y,2);yv=np.load(R/'data/val.npz')['labels'];
assert xt.min()>=0 and xt.max()<=640
mu=xt.mean(0);sd=xt.std(0)+1e-3;z=(xt-mu)/sd;zv=(xv-mu)/sd
A=z.T@z;v=z.T@(yt*2.-1); ym=(yt*2.-1).mean();candidates=[];best=None
for lam in (1,10,100,1000,10000):
 sol=np.linalg.solve(A+lam*np.eye(800),v);w=sol/sd;b=ym-mu@w
 ref=(xv@w+b>0).astype(int);print('lambda',lam,'float',float((ref==yv).mean()),flush=True)
 for limit in (7,15,31):
  scale=limit/abs(w).max();wq=np.rint(w*scale).astype(np.int32);bq=int(round(b*scale));margin=xv.astype(np.int64)@wq+bq
  pred=(margin>0).astype(int);acc=float((pred==yv).mean());row={'lambda':lam,'weight_limit':limit,'validation_accuracy':acc,'ties':int((margin==0).sum()),'absolute_bound':int(640*abs(wq.astype(np.int64)).sum()+abs(bq))};candidates.append(row)
  if row['absolute_bound']>=2**24:continue
  if best is None or acc>best[0]:best=(acc,wq,bq,row)
assert best is not None
_,w,b,row=best;w.tofile(M/'weights.bin');np.array([b],np.int32).tofile(M/'bias.bin')
xv.tofile(R/'validation-counts.bin');np.savez_compressed(R/'training-counts.npz',train=xt,validation=xv)
row=dict(row);row.update(candidates=candidates,classes=['cat','dog'],tie_rule='margin0=>cat; no spike in either output',host='fixed quadrant spike pooling, packet multiplicity only; no learned arithmetic',detector_source='vehicle training only',model_sha256=hashlib.sha256((M/'weights.bin').read_bytes()+(M/'bias.bin').read_bytes()).hexdigest())
(M/'selection.json').write_text(json.dumps(row,indent=2)+'\n');print(json.dumps(row),flush=True)