"""Frozen model evaluation; every learned layer runs in KNP. Exact reference gates."""
from pathlib import Path
import numpy as np,json,subprocess,time,hashlib,csv,argparse
R=Path.home()/'run/pets-knp';M=R/'head-v2';B=Path.home()/'KNP/build/bin'
p=argparse.ArgumentParser();p.add_argument('--stage',choices=['smoke','final'],required=True);a=p.parse_args();O=R/('eval-v2-'+a.stage);O.mkdir(exist_ok=False)
modelhash=hashlib.sha256(b''.join((M/f).read_bytes() for f in ['hidden_w.bin','hidden_b.bin','hidden_thresholds.bin','weights.bin','bias.bin'])).hexdigest()
if a.stage=='final':
 smoke=json.loads((R/'eval-v2-smoke/VERIFIED.json').read_text());assert smoke['cpu_gpu_exact'] and smoke['model_sha256']==modelhash
 (O/'frozen-model.json').write_text(json.dumps({'sha256':modelhash,'selection':json.loads((M/'selection.json').read_text())},indent=2))
def run(cmd,tag):
 start=time.monotonic()
 with (O/(tag+'.log')).open('w') as log:subprocess.run(list(map(str,cmd)),check=True,stdout=log,stderr=subprocess.STDOUT)
 return time.monotonic()-start
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)
def score(raw,labels,tag,cpu=False):
 x=encode(raw);x.tofile(O/(tag+'-encoded.bin'));m=np.load(M/'model.npz');refh=np.clip(np.floor(np.maximum(x.astype(np.int64)@m['w']+m['b'],0)/m['thresholds']),0,64).astype(np.int32);margin=refh.astype(np.int64)@m['readout_w']+int(m['readout_b'][0]);expected=(margin>0).astype(int)
 times={}
 for mode in (['cuda','cpu'] if cpu else ['cuda']):
  hidden=O/(tag+'-'+mode+'-hidden.bin');times[mode+'_hidden']=run([R/'pets_hidden',O/(tag+'-encoded.bin'),M,mode,hidden],tag+'-'+mode+'-hidden');h=np.fromfile(hidden,np.int32).reshape(-1,256);assert np.array_equal(h,refh),(tag,mode,'hidden mismatch',int(np.abs(h-refh).max()))
  out=O/(tag+'-'+mode+'.csv');times[mode+'_readout']=run([R/'pets_readout',hidden,M,mode,out],tag+'-'+mode+'-readout');rows=np.genfromtxt(out,delimiter=',',names=True,dtype=int);assert np.array_equal(rows['prediction'],expected);assert np.array_equal(rows['cat'],(margin<0).astype(int));assert np.array_equal(rows['dog'],(margin>0).astype(int))
 if cpu:assert (O/(tag+'-cuda.csv')).read_bytes()==(O/(tag+'-cpu.csv')).read_bytes()
 return {'total':len(labels),'correct':int((expected==labels).sum()),'accuracy':float((expected==labels).mean()),'confusion':[[int(((labels==c)&(expected==d)).sum()) for d in (0,1)] for c in (0,1)],'ties':int((margin==0).sum()),'hidden_reference_exact':True,'output_reference_exact':True,'cpu_gpu_exact':cpu,'times_seconds':times}
if a.stage=='smoke':
 raw=np.load(R/'training-counts.npz')['validation'][:8];labels=np.load(R/'data/val.npz')['labels'][:8];result=score(raw,labels,'smoke',True)
 # Same sample follows a large input and a repeated sample: hidden reset must give exact counts.
 seq=np.concatenate([raw[:2],raw[1:2],raw[7:8],raw[1:2]]);result['reset_sequence']=score(seq,np.zeros(5,np.int32),'reset',False)
else:
 raw=np.load(R/'training-counts.npz')['validation'];labels=np.load(R/'data/val.npz')['labels'];result={'validation':score(raw,labels,'validation')}
 # Official test first read for final inference after model selection is frozen.
 images=R/'data/test_images.bin';parts=[]
 for first in (0,1000):
  pre=O/f'test-{first}';run([R/'pets_infer',images,first,1000,R/'detectors',B/'libknp-altai-cuda-backend.so',pre],f'test-{first}-detector');parts.append(np.fromfile(str(pre)+'_counts.bin',np.int32).reshape(-1,800));(O/'progress.json').write_text(json.dumps({'current':first+1000,'total':4000,'done':False}))
 raw=np.concatenate(parts);labels=np.load(R/'data/test.npz')['labels'];result['test']=score(raw,labels,'test');(O/'progress.json').write_text(json.dumps({'current':4000,'total':4000,'done':True}));result['test_labels_used_for_selection']=False
assert modelhash==hashlib.sha256(b''.join((M/f).read_bytes() for f in ['hidden_w.bin','hidden_b.bin','hidden_thresholds.bin','weights.bin','bias.bin'])).hexdigest()
result.update(model_sha256=modelhash,stage=a.stage);(O/'VERIFIED.json').write_text(json.dumps(result,indent=2)+'\n');print(json.dumps(result),flush=True)