#!/usr/bin/env python3
"""Sequential full-model ablations; each variant restores the production service."""
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('--variants',nargs='+',default=['attention,shared,dense,head','attention','shared','dense','head'])
p.add_argument('--loss-steps',type=int,default=0)
p.add_argument('--calibrated-w4-families')
p.add_argument('--shared-stage',action='store_true')
a=p.parse_args()
root=Path(os.environ.get('LING3_BOARD_ROOT', '.')).resolve();tools=Path(__file__).resolve().parent
reference=root/'mixed-precision-reference';suite=reference/'suite.json'
a.output.mkdir(parents=True,exist_ok=True)
env=os.environ.copy();env['OPENBLAS_NUM_THREADS']='1';env['OMP_NUM_THREADS']='1'
extra=[]
if a.shared_stage:extra.append('--shared-stage')
if a.calibrated_w4_families:extra.extend(['--calibrated-w4-families',a.calibrated_w4_families])
for families in a.variants:
    output=a.output/families.replace(',','_')
    if output.exists():raise FileExistsError(output)
    print('START '+families,flush=True)
    subprocess.run([sys.executable,str(tools/'validate_official_bridge_board.py'),'--suite',str(suite),
        '--output',str(output),'--only-candidate','--warm-request','--mixed-families',families,
        '--loss-steps',str(a.loss_steps)]+extra,check=True,env=env)
    subprocess.run([sys.executable,str(tools/'score_loss_numpy.py'),'--suite',str(suite),
        '--reference',str(reference),'--candidate',str(output),'--steps',str(a.loss_steps)],check=True,env=env)
    print('COMPLETE '+families,flush=True)
(a.output/'completed.json').write_text(json.dumps({'variants':a.variants,'loss_steps':a.loss_steps,
    'shared_stage':a.shared_stage,'calibrated_w4_families':a.calibrated_w4_families,'status':'complete'},indent=2)+'\n')