#!/usr/bin/env python3
import os
os.environ['OPENBLAS_NUM_THREADS']='1'
from pathlib import Path
import subprocess,concurrent.futures,json,time,signal,sys,hashlib,numpy as np
ROOT=Path('/home/user/run/mnist-altai-native');GROUP=ROOT/(sys.argv[1] if len(sys.argv)>1 else 'gm-shards-v1');GROUP.mkdir(exist_ok=False);START=time.monotonic();N=1000;WORKERS=8
(GROUP/'SHARDS.json').write_text(json.dumps({'total':N,'shards':WORKERS,'replicas':2,'configuration':'default official GM 4x4'}))
children=[];logs=[]
try:
 for i in range(WORKERS):
  ids=np.arange(60000+i*125,60000+(i+1)*125);f=GROUP/f'ids-{i:02}.txt';np.savetxt(f,ids,fmt='%d');name=f'{GROUP.name}/shard-{i:02}';log=(GROUP/f'runner-{i:02}.log').open('w');logs.append(log)
  cmd=[sys.executable,'/home/user/run/altai-official-probe/run_verify.py','gm','125','2',name,str(f)];p=subprocess.Popen(cmd,stdout=log,stderr=subprocess.STDOUT,start_new_session=True);children.append(p)
 (GROUP/'pids.json').write_text(json.dumps([p.pid for p in children]))
 with concurrent.futures.ThreadPoolExecutor(max_workers=WORKERS) as pool:
  future={pool.submit(p.wait,1800):p for p in children}
  for f in concurrent.futures.as_completed(future):
   if f.result()!=0:raise RuntimeError('An official GoldenModel shard failed verification')
 hs=[];outs=[];preds=[];summaries=[]
 for i in range(WORKERS):
  r=GROUP/f'shard-{i:02}';summaries.append(json.loads((r/'VERIFIED.json').read_text()));hs.append(np.fromfile(r/'hidden_counts.bin',np.int32).reshape(125,256));outs.append(np.fromfile(r/'output_counts.bin',np.int32).reshape(125,2,8,10));preds.append(np.loadtxt(r/'predictions.csv',delimiter=',',skiprows=1,dtype=int))
 h=np.concatenate(hs);o=np.concatenate(outs);p=np.concatenate(preds);ref=np.load(ROOT/'reference.npz');assert np.array_equal(p[:,1],np.arange(60000,61000));assert np.array_equal(h,ref['hidden_counts'][:N]);assert np.array_equal(o,ref['output_counts'][:N]);assert np.array_equal(p[:,3],ref['predictions'][:N]);assert np.array_equal(p[:,2],ref['labels'][:N])
 h.tofile(GROUP/'hidden_counts.bin');o.tofile(GROUP/'output_counts.bin');p[:,0]=np.arange(N);np.savetxt(GROUP/'predictions.csv',p,fmt='%d',delimiter=',',header='row,image_id,label,prediction',comments='')
 result=dict(total=N,correct=int((p[:,2]==p[:,3]).sum()),accuracy=float((p[:,2]==p[:,3]).mean()),hidden_exact=True,output_exact=True,predictions_exact=True,requested_ids_exact=True,processes=WORKERS,replicas_per_process=2,grid_per_process='default4x4',backend='official1.0.2 AltAIv1_GoldenModel',wall_seconds=time.monotonic()-START,model_sha256=hashlib.sha256((ROOT/'model.bin').read_bytes()).hexdigest(),shards=summaries)
 (GROUP/'VERIFIED.json').write_text(json.dumps(result,indent=2));print(json.dumps(result),flush=True)
except BaseException as e:
 for p in children:
  if p.poll() is None:
   try:os.killpg(p.pid,signal.SIGTERM)
   except ProcessLookupError:pass
 for p in children:
  try:p.wait(timeout=5)
  except subprocess.TimeoutExpired:os.killpg(p.pid,signal.SIGKILL);p.wait()
 (GROUP/'FAILED.json').write_text(json.dumps({'error':str(e)}));raise
finally:
 for f in logs:f.close()