import datetime
import json
import subprocess
import sys
import time
from pathlib import Path
ROOT=Path('/home/ubuntu/tsr')
OUT=ROOT/'research_round'
PLAN=[{'method':'l2sp','epochs':100,'lr':1e-4,'init':'pretrained','patience':30},
      {'method':'sam','epochs':20,'lr':1e-5,'init':'validation_winner','patience':0},
      {'method':'contrast','epochs':100,'lr':1e-4,'init':'pretrained','patience':30},
      {'method':'consistency','epochs':100,'lr':1e-4,'init':'pretrained','patience':30},
      {'method':'swa','epochs':20,'lr':1e-5,'init':'validation_winner','patience':0}]
(OUT/'plan.json').write_text(json.dumps(PLAN,indent=2))

def command(trial,smoke=False):
    method=trial['method']
    cmd=[sys.executable,'-u',str(ROOT/'finetune_research.py'),'--pretrain_pth',str(ROOT/'latest.utae.tar'),
         '--research',method,'--tag',f'32X32_research_{method}', '--out_dir',str(OUT/('smoke' if smoke else 'checkpoints')/method),
         '--epochs',str(trial['epochs']),'--base_lr',str(trial['lr']),'--encoder_lr_mult','0.5',
         '--batch','2','--crop_size','32','--fixed_crops','--mask_ce_invalid','--resume','none',
         '--patience',str(trial['patience']),'--warmup_epochs', '0' if trial['init']=='validation_winner' else '3']
    if trial['init']=='validation_winner':
        selection=OUT/'initialization.json'
        initial=json.loads(selection.read_text())['checkpoint'] if selection.exists() else str(ROOT/'checkpoints/finetune_v2/h3_mask_ce_invalid/model_best_raw.tar')
        cmd+=['--init_finetuned',initial]
    if not smoke and method == 'contrast':
        cmd+=['--resume','auto']
    if smoke:
        # Exercise consistency at epoch 3 and SWA averaging/BN calibration at epoch 5.
        cmd+=['--epochs',str(6 if method=='swa' else 4 if method=='consistency' else 1),'--limit_batches','1','--patience','0']
    return cmd

def run(trial,smoke=False):
    name=trial['method']; key=('smoke_' if smoke else '')+name
    statusfile=OUT/'queue_status.json'
    status=json.loads(statusfile.read_text()) if statusfile.exists() else {}
    status[key]={'status':'running','started':datetime.datetime.now(datetime.timezone.utc).isoformat(),'command':command(trial,smoke)}
    statusfile.write_text(json.dumps(status,indent=2))
    start=time.monotonic()
    with (OUT/f'{key}.log').open('w') as stream:
        result=subprocess.run(command(trial,smoke),cwd=ROOT,stdout=stream,stderr=subprocess.STDOUT)
    status[key].update(status='completed' if result.returncode==0 else 'failed',seconds=time.monotonic()-start,returncode=result.returncode)
    statusfile.write_text(json.dumps(status,indent=2))
    if result.returncode: raise RuntimeError(f'{key} failed; inspect {key}.log')
    print(key,'completed',status[key]['seconds'],'seconds',flush=True)

def main():
    if '--smoke' in sys.argv:
        for trial in PLAN: run(trial,True)
        return
    import hashlib
    manifest=json.loads((OUT/'protocol_manifest.json').read_text())
    for name,digest in manifest.items():
        if hashlib.sha256((ROOT/name).read_bytes()).hexdigest()!=digest:
            raise RuntimeError(f'Protocol or pretrained file changed: {name}')
    print('Waiting for existing improvement round to finish; no competing GPU training.',flush=True)
    while 'Improvement round complete' not in (ROOT/'improvement_round/service.log').read_text():
        state=subprocess.run(['systemctl','--user','show','tsr-improvement-round.service','--property=ActiveState','--value'],capture_output=True,text=True)
        if state.returncode==0 and state.stdout.strip() in ['failed','inactive']:
            raise RuntimeError('Prior improvement round stopped before completion')
        time.sleep(30)
    (OUT/'initialization.json').write_text((ROOT/'improvement_round/selected_model.json').read_text())
    status=json.loads((OUT/'queue_status.json').read_text())
    for trial in PLAN:
        if status.get('smoke_'+trial['method'],{}).get('status')!='completed': run(trial,True)
    for trial in PLAN:
        current=json.loads((OUT/'queue_status.json').read_text())
        if current.get(trial['method'],{}).get('status') == 'completed':
            print('Retaining completed trial:',trial['method'],flush=True)
            continue
        run(trial)
    print('Research queue completed. Test evaluations run in the separate test service.',flush=True)

if __name__=='__main__': main()