"""Load the saved best FNO checkpoint and evaluate untouched official test data."""
from pathlib import Path

from common import (case_guard, config_hash, load_config, load_data,
                    negative_total_flow_rates, parse_args, predict, save_json,
                    save_torch, sha256)
import torch
from model.NOTS import NOTS


def main():
    parser = parse_args(__doc__)
    parser.add_argument('--checkpoint', type=Path)
    args = parser.parse_args()
    cfg = load_config(args.config)
    device = case_guard(cfg)
    out = Path(cfg['paths']['output_dir'])
    checkpoint_path = args.checkpoint or out / 'weight/best_model.pt'
    state = torch.load(checkpoint_path, weights_only=False, map_location='cpu')
    if state['config_hash'] != config_hash(cfg):
        raise ValueError('Checkpoint/config protocol mismatch')
    if state['normalization']['kind'] != 'identity':
        raise ValueError('Unexpected checkpoint prediction scale')
    model = NOTS(**state['model_config']).to(device)
    model.load_state_dict(state['model_state'], strict=True)
    model.eval()
    metrics = {'checkpoint': str(checkpoint_path), 'checkpoint_sha256': sha256(checkpoint_path),
               'selected_seed': state['seed'], 'selection_train_mse': state['train_mse'], 'splits': {}}
    for split in ('test', 'train'):
        x, y, meta = load_data(cfg, split)
        if split == 'train' and meta['sha256'] != state['data_sha256']:
            raise ValueError('Checkpoint training pool hash differs')
        x, y = x.to(device), y.to(device)
        with torch.no_grad():
            output = predict(model, x, cfg['experiment']['prediction_batch_size'])
            error = output - y
            truth = negative_total_flow_rates(y, x)
            predicted_flow = negative_total_flow_rates(output, x)
            result = {'n': len(x), 'mse': error.square().mean().item(),
                      'relative_l2_mean': (error.flatten(1).norm(dim=1) / y.flatten(1).norm(dim=1).clamp_min(1e-12)).mean().item(),
                      'functional_mae': (truth - predicted_flow).abs().mean().item(),
                      'predicted_best_index': predicted_flow.argmax().item(),
                      'true_best_index': truth.argmax().item(),
                      'prediction_shape': list(output.shape),
                      'all_finite': bool(torch.isfinite(output).all().item()), 'data': meta}
        if not result['all_finite']:
            raise FloatingPointError('Nonfinite checkpoint inference')
        metrics['splits'][split] = result
        save_torch(out / 'results' / f'inference_{split}.pt', {'x': x.cpu(), 'target': y.cpu(),
                   'prediction': output.cpu(), 'true_functional': truth.cpu(),
                   'predicted_functional': predicted_flow.cpu(), 'checkpoint_sha256': metrics['checkpoint_sha256']})
        print(f'INFERENCE split={split} {result}', flush=True)
    save_json(out / 'results/inference_metrics.json', metrics)


if __name__ == '__main__':
    main()