#!/usr/bin/env python3
"""Five warmed requests per candidate with the same shared-input optimization."""
import argparse
import json
import os
from pathlib import Path
import subprocess
import sys

p=argparse.ArgumentParser(description=__doc__)
p.add_argument('--output',type=Path,required=True)
p.add_argument('--resume',action='store_true',help='Skip fully completed matching jobs after a preflight failure')
a=p.parse_args()
script=Path(__file__).resolve().parent/'validate_official_bridge_board.py'
suite=Path(os.environ.get('LING3_BOARD_SUITE', 'mixed-precision-reference/suite.json')).resolve()
a.output.mkdir(parents=True,exist_ok=True)
jobs=[('scheme3_shared_stage','attention,shared',[]),
      ('selective_calibrated','kda_out,mla_in,mla_out,shared,dense,head',['--calibrated-w4-families','kda_in'])]
for name,families,extra in jobs:
    directory=a.output/name
    if directory.exists():
        if not a.resume:raise FileExistsError(directory)
        restored=json.loads((directory/'health-restored.json').read_text())
        perf=json.loads((directory/('mixed_'+families.replace(',','_')+'-performance.json')).read_text())
        expected_cal='kda_in' if extra else None
        if not (restored['status']=='ok' and restored['context_length']==65536 and
                len(perf['runs'])==5 and perf['summary']['exact_128_64'] and
                perf['configuration']['w8_families']==families and perf['configuration']['shared_stage'] and
                perf['configuration']['calibrated_w4_families']==expected_cal):
            raise RuntimeError('incomplete or mismatched prior job: '+name)
        print('REUSE completed '+name,flush=True)
        continue
    print('START '+name,flush=True)
    subprocess.run([sys.executable,str(script),'--suite',str(suite),'--output',str(directory),
        '--only-candidate','--skip-loss','--warm-request','--runs','5','--shared-stage',
        '--mixed-families',families]+extra,check=True)
    print('COMPLETE '+name,flush=True)
(a.output/'completed.json').write_text(json.dumps({'status':'complete','jobs':[j[0] for j in jobs]},indent=2)+'\n')