"""Run the eight fixed phase-budget arms with at most two workers."""
from concurrent.futures import ThreadPoolExecutor
from datetime import datetime, timezone
import hashlib
import json
import os
from pathlib import Path
import subprocess
import sys
root=Path(__file__).resolve().parent
python='/workspace/cadence-robot-arena/.venv/bin/python'
env={**os.environ,'OPENBLAS_NUM_THREADS':'1','OMP_NUM_THREADS':'1','MKL_NUM_THREADS':'1'}
protocol=json.loads((root/'protocol.json').read_text())
admission=json.loads((root/'sources-control.json').read_text())
reused=json.loads((root/'reused-controls.json').read_text())
for case,record in reused.items():
 data=Path(record['receipt']).read_bytes()
 assert hashlib.sha256(data).hexdigest()==record['receipt_sha256']==protocol['reused_controls_sha256'][case]
 receipt=json.loads(data)
 assert receipt['status']=='complete' and receipt['source_unchanged'] and receipt['producer_unchanged']
 assert receipt['source_sha256']==admission['source_sha256']
 assert receipt['training']['total']['moments']==protocol['additional_training_moments']
 assert {(r['policy'],r['seed'],r['total']['moments']) for r in receipt['evaluation']}=={(name,seed,protocol['evaluation_moments']) for name in ['trained','ablated','newborn','random'] for seed in protocol['evaluation_worlds']}
jobs=[(case,'full-budget') for case in protocol['first_stage']]+[(case,'control') for case in ['hexapod-302','hexapod-301']]
manifest={'started_utc':datetime.now(timezone.utc).isoformat(),'jobs':jobs,'max_workers':2,'hashes':{name:hashlib.sha256((root/name).read_bytes()).hexdigest() for name in ['compare.py','protocol.json','sources-control.json','sources-full-budget.json','reused-controls.json','launch.py']}}
(root/'launch.json').write_text(json.dumps(manifest,indent=2)+'\n')
def run(job):
 case,arm=job
 command=[python,str(root/'compare.py'),'--case',case,'--arm',arm]
 with (root/f'{case}-{arm}.log').open('w') as stream:
  result=subprocess.run(command,env=env,stdout=stream,stderr=subprocess.STDOUT)
 out={'case':case,'arm':arm,'returncode':result.returncode,'finished_utc':datetime.now(timezone.utc).isoformat()}
 print(json.dumps(out),flush=True)
 return out
with ThreadPoolExecutor(max_workers=2) as pool:
 results=list(pool.map(run,jobs))
(root/'complete.json').write_text(json.dumps(results,indent=2)+'\n')
sys.exit(any(row['returncode'] for row in results))
