"""Freeze an owned training checkpoint and evaluate only development splits."""
import argparse
import json
import os
from pathlib import Path
import subprocess
import sys
import torch
from safetensors.torch import save_file


def main():
    p=argparse.ArgumentParser();p.add_argument('--checkpoint',required=True);p.add_argument('--data',required=True)
    p.add_argument('--out',required=True);args=p.parse_args();data=Path(args.data)
    r=torch.load(args.checkpoint,map_location='cpu',weights_only=False)
    dest=Path(args.out)/('step-'+str(r['step']));dest.mkdir(parents=True,exist_ok=True)
    if (dest/'selection.json').exists():print((dest/'selection.json').read_text());return
    save_file({k:v.to(torch.bfloat16).contiguous() for k,v in r['model'].items()},str(dest/'model.safetensors'),
              metadata={'step':str(r['step']),'random_initialization':'true'})
    (dest/'config.json').write_text(json.dumps(r['config'],indent=2))
    info={k:r[k] for k in ['step','processed_tokens','response_tokens','training_seconds','random_initialization']}
    (dest/'checkpoint-info.json').write_text(json.dumps(info,indent=2));del r
    print(json.dumps({'frozen':str(dest),'step':info['step']}),flush=True);scores={};digest=None
    for split in ['validation','development']:
        if not (data/(split+'.jsonl')).exists():continue
        with (dest/(split+'.log')).open('w') as log:
            subprocess.run([sys.executable,'-u','-m','tinyquery.evaluate','--checkpoint',str(dest/'model.safetensors'),
                            '--tokenizer',str(data/'tokenizer.json'),'--data',str(data/(split+'.jsonl')),
                            '--out',str(dest/(split+'.jsonl'))],stdout=log,stderr=subprocess.STDOUT,check=True)
        summary=json.loads((dest/(split+'.summary.json')).read_text());scores[split]=summary['metrics']['success']
        digest=summary['checkpoint_sha256'];print(json.dumps({'split':split,'metrics':summary['metrics']}),flush=True)
    selection={'checkpoint':str(dest/'model.safetensors'),'step':info['step'],'checkpoint_sha256':digest,
               'development_scores':scores,'selection_score':sum(v['rate'] for v in scores.values())/len(scores),
               'rule':'Equal mean of full validation and development task success. Test/manual excluded.'}
    (dest/'selection.json').write_text(json.dumps(selection,indent=2));print(json.dumps(selection),flush=True)


if __name__=='__main__':main()