"""Verify full-budget completion and checkpoint portability on the compute node."""
import csv
import json
import sys
from pathlib import Path
sys.path.insert(0,str(Path(__file__).resolve().parent))
import numpy as np
import torch
import yaml
from common import arguments,config,path,write_json
from inference import run as infer


def run(cfg):
    result=path(cfg['paths']['results'])
    history=list(csv.DictReader((result/'history.csv').open()))
    expected=cfg['training']['epochs']
    assert expected==10000 and len(history)==expected
    assert [int(r['epoch']) for r in history]==list(range(1,expected+1))
    assert int(history[-1]['global_step'])==160000
    assert all(np.isfinite(float(row[key])) for row in history for key in row)
    summary=json.loads((result/'training_summary.json').read_text())
    assert summary['status']=='completed' and not summary['early_stopping']
    ck=torch.load(path(cfg['paths']['best_checkpoint']),map_location='cpu',weights_only=False)
    best=min(history,key=lambda r:float(r['validation_relative_l2_normalized_batch_mean']))
    assert ck['epoch']==int(best['epoch'])
    assert ck['best_validation']==float(best['validation_relative_l2_normalized_batch_mean'])
    assert sum(v.numel() for v in ck['model_state_dict'].values())==101100
    original=np.load(result/'predictions.npz',allow_pickle=False)
    assert original['prediction'].shape==(100,500)
    assert np.isfinite(original['prediction']).all()
    assert np.array_equal(original['original_indices'],np.arange(900,1000))
    # Repoint data AND result directories and reload the same checkpoint.
    alias=result/'portability_data'
    target=path(cfg['data']['directory']).resolve()
    if not alias.exists():alias.symlink_to(target,target_is_directory=True)
    portable=json.loads(json.dumps(cfg))
    portable['data']['directory']=str(alias)
    portable['paths']['results']=str(result/'portability_check')
    (result/'portability_check').mkdir(exist_ok=True)
    (result/'portability_check/config.yaml').write_text(yaml.safe_dump(portable,sort_keys=False))
    infer(portable)
    repeated=np.load(result/'portability_check/predictions.npz',allow_pickle=False)
    assert np.array_equal(original['prediction'],repeated['prediction'])
    report=dict(status='passed',formal_epochs=10000,optimizer_updates=160000,batch_size=50,
        sensors=200,output_points=500,parameter_count=101100,best_epoch=ck['epoch'],
        best_checkpoint_selected_by_validation_only=True,test_indices='900:1000',
        portability=dict(data_directory_changed=True,results_directory_changed=True,
                         same_checkpoint=True,predictions_bitwise_equal=True),
        checkpoint_normalization_self_contained=True)
    write_json(result/'verification.json',report)
    print('VERIFICATION_COMPLETE '+json.dumps(report),flush=True)


if __name__=='__main__':
    a=arguments(__doc__).parse_args();run(config(a.config))