#!/usr/bin/env python3
import os
os.environ['OPENBLAS_NUM_THREADS']='2'
from pathlib import Path
import json,subprocess,sys,time,hashlib,numpy as np
root=Path('/home/user/run/mnist-altai-native');mode=sys.argv[1];n=int(sys.argv[2]);rep=int(sys.argv[3]);name=sys.argv[4];idsfile=sys.argv[5] if len(sys.argv)>5 else '-';out=root/name
if out.exists() and any(out.iterdir()):raise SystemExit('Refusing existing nonempty output directory: '+str(out))
out.mkdir(exist_ok=True)
cmd=['/home/user/run/altai-official-probe/native_digits',str(root),str(out),'60000',str(n),str(rep),mode,idsfile]
(out/'command.json').write_text(json.dumps(cmd));start=time.monotonic();model_hash=hashlib.sha256((root/'model.bin').read_bytes()).hexdigest();exe_hash=hashlib.sha256(Path(cmd[0]).read_bytes()).hexdigest()
expected_ids=np.arange(60000,60000+n) if idsfile=='-' else np.loadtxt(idsfile,dtype=int).reshape(-1)[:n]
if len(expected_ids)!=n:raise SystemExit('Insufficient requested image IDs')
try:
 with (root/(name+'.log')).open('w') as log:
  child=subprocess.Popen(cmd,stdout=log,stderr=subprocess.STDOUT,cwd=out);(out/'child.pid').write_text(str(child.pid));reference=np.load(root/'reference.npz')
  while True:
   try:rc=child.wait(timeout=5);break
   except subprocess.TimeoutExpired:
    if time.monotonic()-start>1800:raise RuntimeError('Native run exceeded 1800s')
    if (out/'hidden_counts.bin').exists() and (out/'output_counts.bin').exists():
     completed=min((out/'hidden_counts.bin').stat().st_size//(256*4),(out/'output_counts.bin').stat().st_size//(160*4))
     if completed and idsfile=='-':
      hh=np.fromfile(out/'hidden_counts.bin',np.int32,count=completed*256).reshape(completed,256);oo=np.fromfile(out/'output_counts.bin',np.int32,count=completed*160).reshape(completed,2,8,10)
      assert np.array_equal(hh,reference['hidden_counts'][:completed]),'incremental hidden counter mismatch'
      assert np.array_equal(oo,reference['output_counts'][:completed]),'incremental output counter mismatch'
 if rc:raise RuntimeError(f'native process exit {rc}')
 rows=np.loadtxt(out/'predictions.csv',delimiter=',',skiprows=1,dtype=int).reshape(n,4);ref=np.load(root/'reference.npz');ids=rows[:,1]-60000
 if not np.array_equal(rows[:,0],np.arange(n)) or not np.array_equal(rows[:,1],expected_ids) or (ids<0).any() or (ids>=10000).any():raise RuntimeError('invalid ids/order')
 h=np.fromfile(out/'hidden_counts.bin',np.int32).reshape(n,256);o=np.fromfile(out/'output_counts.bin',np.int32).reshape(n,2,8,10)
 assert np.array_equal(h,ref['hidden_counts'][ids]),f'hidden differs max={abs(h-ref["hidden_counts"][ids]).max()}'
 assert np.array_equal(o,ref['output_counts'][ids]),f'output differs max={abs(o-ref["output_counts"][ids]).max()}'
 assert np.array_equal(rows[:,2],ref['labels'][ids]);assert np.array_equal(rows[:,3],ref['predictions'][ids])
 assert hashlib.sha256((root/'model.bin').read_bytes()).hexdigest()==model_hash,'Model changed during run'
 assert hashlib.sha256(Path(cmd[0]).read_bytes()).hexdigest()==exe_hash,'Executable changed during run'
 summary=json.loads((out/'summary.json').read_text());summary.update(requested_ids_exact=True,executable_sha256=exe_hash,hidden_exact=True,output_exact=True,predictions_exact=True,model_sha256=hashlib.sha256((root/'model.bin').read_bytes()).hexdigest(),runner_wall_seconds=time.monotonic()-start)
 (out/'VERIFIED.json').write_text(json.dumps(summary,indent=2));print(json.dumps(summary),flush=True)
except Exception as e:
 if 'child' in locals() and child.poll() is None:child.kill();child.wait()
 (out/'FAILED.json').write_text(json.dumps({'error':str(e)}));raise