#!/usr/bin/env python3
"""Pin every complete production checkpoint through the requested stage."""
import argparse
import hashlib
import json
from pathlib import Path

HERE = Path(__file__).resolve().parent
SC = Path('/e/scratch/reformo/schuhmann1_moss')
CHECKPOINTS = SC / 'out/m2_600m_ladder/caption_curriculum_v1/training/checkpoints'
PLAN = HERE / 'plan_caption_production_32nodes.json'
PROMPTS = SC / 'out/m2_20k_score/eval/timed_prompts_v3.json'


def sha(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open('rb') as stream:
        for block in iter(lambda: stream.read(8 << 20), b''):
            digest.update(block)
    return digest.hexdigest()


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--through-stage', default='S5')
    ap.add_argument('--output', default=str(HERE / 'validation_sweep_S1_S5_all_checkpoints.json'))
    args = ap.parse_args()
    plan = json.loads(PLAN.read_text())
    boundaries, end = [], 0
    for stage in plan['stages']:
        end += int(stage['updates'])
        boundaries.append((stage['name'], end))
    assert args.through_stage in dict(boundaries)
    final_step = dict(boundaries)[args.through_stage]
    candidates = []
    for directory in sorted(CHECKPOINTS.glob('step-*')):
        complete = directory / 'COMPLETE.json'
        if not complete.exists():
            continue
        proof = json.loads(complete.read_text())
        step = int(proof['step'])
        if step > final_step:
            continue
        assert proof['world'] == plan['world_size']
        assert proof['contract']['plan_sha256'] == sha(PLAN)
        assert proof['contract']['objective'] == 'global_token_mean_continuous_ladder_v1'
        stage, _ = next((name, bound) for name, bound in boundaries if step <= bound)
        state = directory / 'training_state.pt'
        assert state.stat().st_size == int(proof['state_bytes'])
        candidates.append({'stage': stage, 'step': step, 'checkpoint': str(directory),
                           'state_sha256': proof['state_sha256'],
                           'export_available': (directory / 'model_bf16.pt').exists()})
    assert candidates and max(row['step'] for row in candidates) == final_step
    stage_order = {name: index for index, (name, _) in enumerate(boundaries)}
    candidates.sort(key=lambda row: (-stage_order[row['stage']], row['step']))
    payload = {'version': 1, 'purpose': 'all complete checkpoints through selected stage',
        'metric': 'independent_heldout_validation_loss',
        'selection_policy': 'minimum pooled token-weighted heldout loss per stage',
        'acting_metrics_used_for_selection': False, 'training_loss_used_for_selection': False,
        'world': 4, 'limit': None, 'production_plan': str(PLAN),
        'production_plan_sha256': sha(PLAN), 'heldout_prompts': str(PROMPTS),
        'heldout_prompts_sha256': sha(PROMPTS), 'heldout_rows': 50,
        'tasks_per_checkpoint': 100, 'through_stage': args.through_stage,
        'candidate_count': len(candidates), 'candidates': candidates}
    output = Path(args.output).resolve()
    assert not output.exists(), f'refusing to overwrite existing manifest: {output}'
    output.write_text(json.dumps(payload, indent=2) + '\n')
    print(json.dumps({'manifest': str(output), 'candidate_count': len(candidates),
                      'stages': {name: sum(c['stage'] == name for c in candidates)
                                 for name, bound in boundaries if bound <= final_step}}, indent=2))


if __name__ == '__main__':
    main()