#!/usr/bin/env python3
"""Evaluate best and final trained checkpoints at every original SU2 point."""
import argparse
import json
import os
from pathlib import Path
import sys

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
import numpy as np
import torch
from scripts.common import load_config, build_model, read_reference, predict, errors, atomic_json, require_compute_node


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--config', default='config/config.yaml')
    parser.add_argument('--checkpoint', type=Path)
    args = parser.parse_args()
    cfg, root = load_config(args.config)
    require_compute_node(cfg)
    device = torch.device(cfg['training']['device'])
    results = root / cfg['paths']['results']
    results.mkdir(parents=True, exist_ok=True)
    xy, ref, meta, mesh = read_reference(cfg['evaluation']['reference_vtu'], cfg['evaluation']['expected_points'])
    paths = [('best', root / cfg['paths']['best_checkpoint']), ('final', root / cfg['paths']['last_checkpoint'])]
    if args.checkpoint:
        paths = [('custom', args.checkpoint)]
    metrics = {'reference': meta, 'definition': 'norm(predicted Velocity.x - SU2 Velocity.x) / norm(SU2 Velocity.x)',
               'all_original_nodes': True, 'volume_weighting': False, 'checkpoints': {}}
    for label, path in paths:
        ck = torch.load(path, map_location=device, weights_only=False)
        if ck['step'] < cfg['training']['steps'] and label == 'final':
            raise RuntimeError(f'Formal final checkpoint contains only {ck["step"]} steps')
        model = build_model(ck['config'], device)
        model.load_state_dict(ck['model_state_dict'], strict=True)
        model.eval()
        pred = predict(model, xy, device, cfg['evaluation']['batch_size'])
        if not np.isfinite(pred).all():
            raise FloatingPointError('Nonfinite inference prediction')
        result = dict(checkpoint=str(path), checkpoint_step=ck['step'], **errors(pred, ref))
        metrics['checkpoints'][label] = result
        np.savez_compressed(results / f'predictions_{label}.npz', xy=xy, prediction_uvp=pred,
                            reference_uvp=ref, point_indices=np.arange(len(xy)), checkpoint_step=ck['step'])
        from vtk import vtkXMLUnstructuredGridWriter
        from vtk.util.numpy_support import numpy_to_vtk
        for name, array in [('PINN_u', pred[:, 0]), ('PINN_v', pred[:, 1]),
                            ('PINN_p', pred[:, 2]), ('absolute_error_u', abs(pred[:, 0] - ref[:, 0]))]:
            a = numpy_to_vtk(np.ascontiguousarray(array), deep=True)
            a.SetName(name)
            mesh.GetPointData().AddArray(a)
        writer = vtkXMLUnstructuredGridWriter()
        writer.SetFileName(str(results / f'prediction_{label}.vtu'))
        writer.SetInputData(mesh)
        if writer.Write() != 1:
            raise IOError('VTU prediction export failed')
        print(json.dumps({'event': 'inference', 'selection': label, **result}), flush=True)
    metrics['paper_relative_l2_u'] = cfg['evaluation']['paper_relative_l2_u']
    metrics['slurm_job_id'] = os.environ.get('SLURM_JOB_ID')
    atomic_json(results / 'metrics.json', metrics)


if __name__ == '__main__':
    main()