import json, subprocess, sys, time
from pathlib import Path
ROOT=Path('/home/ubuntu/tsr')
OUT=ROOT/'research_round'
PLAN=json.loads((OUT/'plan.json').read_text())
while 'Research queue completed.' not in (OUT/'recovery_service.log').read_text() and 'Research queue and five' not in (OUT/'recovery_service.log').read_text():
    time.sleep(30)
def evaluate_all():
    # Select one checkpoint per method by fixed-crop validation, then test once.
    for trial in PLAN:
        method=trial['method']
        folder=OUT/'checkpoints'/method
        history=json.loads((folder/'history.json').read_text())
        candidates=[(row[kind]['miou'],kind,row['epoch']+1)
                    for row in history for kind in ['raw','ema','swa'] if row.get(kind)]
        score,kind,epoch=max(candidates)
        checkpoint=folder/f'model_best_{kind}.tar'
        selection={'method':method,'kind':kind,'epoch':epoch,'validation_miou':score,'checkpoint':str(checkpoint)}
        (folder/'test_selection.json').write_text(json.dumps(selection,indent=2))
        result=folder/'test_results.json'
        if not result.exists():
            print('Testing validation-selected checkpoint:',selection,flush=True)
            with (OUT/f'{method}_test.log').open('w') as stream:
                subprocess.run([sys.executable,'-u',str(ROOT/'evaluate_test.py'),'--ckpt',str(checkpoint),
                                '--batch_size','1','--output',str(result)],cwd=ROOT,stdout=stream,
                               stderr=subprocess.STDOUT,check=True)

evaluate_all()
print('All five research test evaluations completed',flush=True)