"""One process per CPU/GPU trial, gate the ordered fixed protocol."""
import argparse
import json
import os
import sys
import time
import traceback
from context import ROOT,MODEL,EXP,CFG,check,save,budget_check

def main():
    ap=argparse.ArgumentParser();ap.add_argument('backend',choices=['oracle','gpu']);ap.add_argument('case')
    a=ap.parse_args();check();budget_check()
    names=[x['name'] for x in CFG['trials']];assert a.case in names
    position=names.index(a.case)
    if position:
        previous=json.loads((EXP/'audit'/('completed-'+names[position-1]+'.json')).read_text())
        assert previous['pass_all'], 'Previous fixed trial failed; do not continue'
    case=EXP/'cases'/a.case;out=EXP/a.backend/a.case
    assert not out.exists(), 'Never overwrite an executed trial'
    start=time.perf_counter()
    try:
        if a.backend=='oracle':
            from oracle_full import run_case
            result=run_case(case,out)
        else:
            os.environ['CUDA_VISIBLE_DEVICES']=CFG['gpu_uuid'];os.environ['CUDA_HOME']='/usr/local/cuda'
            os.environ['PATH']=str(ROOT/'models/male-cns-lif/.venv/bin')+':/usr/local/cuda/bin:'+os.environ['PATH']
            os.environ['MAX_JOBS']='2';os.environ['TORCH_CUDA_ARCH_LIST']='8.9'
            sys.path.insert(0,str(ROOT/'models/male-cns-lif/terminal-scope-diagnostic-v1'))
            from diagnostic import queues
            queues()
            from gpu_full import run_gpu
            result=run_gpu(case,out,EXP/'oracle'/a.case,EXP/'extension-build')
        save(EXP/'audit'/(a.backend+'-'+a.case+'.json'),dict(case=a.case,backend=a.backend,
             wall_seconds_including_validation=time.perf_counter()-start,result=result))
        budget_check();print('COMPLETED',a.backend,a.case,json.dumps(result),flush=True)
    except Exception as exc:
        save(EXP/'audit'/('error-'+a.backend+'-'+a.case+'.json'),dict(error=str(exc),type=type(exc).__name__,traceback=traceback.format_exc()))
        raise

if __name__=='__main__':main()