#!/usr/bin/env python3
import os
for name in ['OPENBLAS_NUM_THREADS','OMP_NUM_THREADS','MKL_NUM_THREADS']:os.environ[name]='8'
import base64,hashlib,json,time,traceback
from pathlib import Path
import numpy as np
from model_core import normalize,augment,predict,quantized,CAP,SCALE
ROOT=Path('/home/user/run/mnist-fast');ROOT.mkdir(exist_ok=True,parents=True)
START=time.time(); state={'current':0,'total':24,'done':False,'failed':False,'label':'normalize MNIST'}
def dump(path,obj):
 p=path.with_suffix(path.suffix+'.tmp');p.write_text(json.dumps(obj,separators=(',',':'))+'\n');p.replace(path)
def progress(**values):state.update(values,wall_seconds=time.time()-START);dump(ROOT/'progress.json',state)
def main():
 progress();rng=np.random.default_rng(20261015)
 data=Path('/home/user/run/mnist-native-20261015');raw=np.fromfile(data/'MNIST.bin',np.uint8).reshape(-1,28,28);labels=np.array([int(x) for x in (data/'MNIST.target').read_text().split()],np.int64)
 assert raw.shape==(70000,28,28)
 cache=ROOT/'normalized.npy'
 if cache.exists():x=np.load(cache)
 else:
  x=np.empty((70000,784),np.uint8)
  for first in range(0,70000,1000):
   x[first:first+1000]=np.stack([normalize(im).ravel() for im in raw[first:first+1000]]);progress(current=(first+1000)/70000,label='normalize bbox20 + center of mass')
  np.save(cache,x)
 order=rng.permutation(60000);train=order[:55000];val=order[55000:];np.savez(ROOT/'split.npz',train=train,validation=val,test=np.arange(60000,70000))
 w1=rng.normal(0,np.sqrt(2/784),(784,256)).astype(np.float32);b1=np.full(256,.05,np.float32);w2=rng.normal(0,np.sqrt(2/256),(256,10)).astype(np.float32);b2=np.zeros(10,np.float32)
 params=[w1,b1,w2,b2];mom=[np.zeros_like(p) for p in params];vel=[np.zeros_like(p) for p in params];updates=0;best=-1;history=[]
 xv=x[val].astype(np.float32)/255;yv=labels[val]
 for epoch in range(20):
  shuffled=rng.permutation(train);loss=0;lr=.001*(.35 if epoch>=12 else 1)*(.35 if epoch>=17 else 1)
  for offset in range(0,len(shuffled),256):
   ids=shuffled[offset:offset+256];a=augment(x[ids],rng);y=labels[ids];z=a@w1+b1;h=np.clip(z,0,CAP)
   # Surrogategradient trainingthroughquantizedIFcounts, switchedonafterwarmup.
   hq=np.floor(h*(64/CAP))*(CAP/64) if epoch>=2 else h
   logits=hq@w2+b2;logits-=logits.max(1,keepdims=True);e=np.exp(logits);prob=e/e.sum(1,keepdims=True);loss+=float(-np.log(np.maximum(prob[np.arange(len(ids)),y],1e-12)).sum())
   dz=prob;dz[np.arange(len(ids)),y]-=1;dz/=len(ids)
   g2=hq.T@dz+.00005*w2;gb2=dz.sum(0);dh=(dz@w2.T)*((z>0)&(z<CAP));g1=a.T@dh+.00005*w1;gb1=dh.sum(0)
   updates+=1
   for p,g,m,v in zip(params,[g1,gb1,g2,gb2],mom,vel):
    m*=.9;m+=.1*g;v*=.999;v+=.001*g*g;p-=lr*(m/(1-.9**updates))/(np.sqrt(v/(1-.999**updates))+1e-8)
  pred=predict(x[val],params,64);acc=float((pred==yv).mean());history.append({'epoch':epoch+1,'validation_accuracy':acc,'train_aug_loss':loss/55000,'elapsed':time.time()-START});print(history[-1],flush=True)
  if acc>best:
   best=acc;np.savez(ROOT/'best.npz',w1=w1,b1=b1,w2=w2,b2=b2)
  dump(ROOT/'history.json',history);progress(current=1+epoch+1,label=f'train epoch{epoch+1}/20; validation{acc:.4f}')
 best=np.load(ROOT/'best.npz');params=[best[k] for k in ['w1','b1','w2','b2']]
 candidates=[]
 for steps in [32,64,128]:
  pred=predict(x[val],params,steps);candidates.append({'steps':steps,'accuracy':float((pred==yv).mean())})
 chosen=max(candidates,key=lambda d:(d['accuracy'],-d['steps']));steps=chosen['steps'];dump(ROOT/'selection.json',{'method':'checkpoint and IFsteps selected only on5000validation','validation_candidates':candidates,'chosen':chosen,'architecture':'784 currentinputs →256integerIFhidden →10analogspike-countreadout','seed':20261015,'train':55000,'validation':5000,'test':10000})
 progress(current=22,label='fixed model; evaluate officialtest10000 and stress')
 test=x[60000:];ytest=labels[60000:];scores=predict(test,params,steps,return_scores=True);pred=scores.argmax(1);correct=int((pred==ytest).sum());conf=np.zeros((10,10),int);np.add.at(conf,(ytest,pred),1)
 # Fixed stressprotocol afterselection; no tuningfromtheseresults.
 stress={}
 for name,operation in [('shift_right2',None),('thicker_cross1','max'),('thinner_cross1','min')]:
  a=test.reshape(-1,28,28)
  if operation is None:
   transformed=np.zeros_like(a);transformed[:,:,2:]=a[:,:,:-2]
  else:
   p=np.pad(a,((0,0),(1,1),(1,1)));op=np.maximum if operation=='max' else np.minimum
   transformed=op.reduce([p[:,1:29,1:29],p[:,:28,1:29],p[:,2:30,1:29],p[:,1:29,:28],p[:,1:29,2:30]])
  stress[name]=float((predict(transformed.reshape(-1,784),params,steps)==ytest).mean())
  progress(current=22+len(stress)/3,label='evaluate fixed stress protocol')
 # Separate outputstateandtestregression; screenshotindexknown, neverselectedon.
 regression={'test_index':2534,'label':int(ytest[2534]),'prediction':int(pred[2534]),'used_for_training_or_selection':False}
 baseline=np.loadtxt('/home/user/run/mnist-native-20261015/predictions.csv',delimiter=',',skiprows=1,dtype=int);assert len(baseline)==10000 and np.array_equal(baseline[:,1],ytest)
 metrics={'test_total':10000,'test_correct':correct,'test_accuracy':correct/10000,'baseline_same_test_correct':int((baseline[:,2]==ytest).sum()),'baseline_same_test_accuracy':float((baseline[:,2]==ytest).mean()),'validation_selection':chosen,'steps':steps,'confusion':conf.tolist(),'stress_postselection':stress,'regression':regression,'wall_seconds':time.time()-START,'knp_verified':False,'output':'analoglinearreadoutofhiddenIFspikecounts','training':'offlineANNsurrogatequantization; notSTDP','preprocessing':'bbox20 + integerCOMshift; codeidenticalPythonJS'}
 dump(ROOT/'metrics.json',metrics);np.savez(ROOT/'test_results.npz',predictions=pred,labels=ytest,scores=scores)
 q=quantized(params);model={'version':'mnist-if256-20261015','architecture':'784 current pixels →256 integer integrate-and-fire neurons →10 analog spike-count scores','knp_verified':False,'steps':steps,'cap':CAP,'scale':SCALE,'threshold':int(CAP*SCALE*255),'metrics':metrics,'shape':[784,256,10],'weights':{k:base64.b64encode(a.astype('<i2' if a.dtype==np.int16 else '<i4').tobytes()).decode() for k,a in zip(['w1','b1','w2','b2'],q)}}
 dump(ROOT/'model.json',model)
 # Gallery contains only officialtest, withprovenance,index2534regressionincluded.
 indices=[2534,0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28]
 dump(ROOT/'gallery.json',{'split':'official MNIST test; nevertrainedorselectedon','items':[{'index':i,'label':int(ytest[i]),'pixels':base64.b64encode(raw[60000+i].tobytes()).decode()} for i in indices]})
 dump(ROOT/'provenance.json',{'dataset_manifest':json.loads((data/'dataset-manifest.json').read_text()),'split_sha256':hashlib.sha256((ROOT/'split.npz').read_bytes()).hexdigest(),'model_sha256':hashlib.sha256((ROOT/'model.json').read_bytes()).hexdigest(),'training_seed':20261015,'normalization':'see model_core.py','no_test_tuning':True})
 progress(current=24,done=True,label=f'finished; officialtest{correct}/10000');print(json.dumps(metrics),flush=True)
if __name__=='__main__':
 try:main()
 except Exception as e:progress(failed=True,error=str(e),label='failed');traceback.print_exc();raise