"""Evaluate all 200 Darcy examples with the independently loaded best weights."""
from __future__ import annotations

import csv
from pathlib import Path
import time

import numpy as np
import torch
from torch.utils.data import DataLoader, TensorDataset

from common import (atomic_json, compute_device, configuration, decode_output,
                    encode_input, load_checkpoint, read_config, read_split, relative_l2)


def main():
    parser=configuration('HNO Darcy best-checkpoint inference')
    parser.add_argument('--checkpoint', help='Override the checkpoint path in config')
    args=parser.parse_args()
    c=read_config(args.config)
    device=compute_device(args.device)
    torch.set_num_threads(c['training']['cpu_threads'])
    torch.backends.cudnn.benchmark=True
    checkpoint_path=Path(args.checkpoint or c['output']['checkpoint'])
    model,stats,checkpoint=load_checkpoint(checkpoint_path,device)
    if checkpoint['model_config']!=c['model']:
        raise ValueError('Current model config disagrees with saved architecture')
    model.eval()
    print(f'Loaded best weights: {checkpoint_path}; best epoch={checkpoint["epoch"]}; '
          f'selection relative L2={checkpoint["eval_relative_l2"]:.9f}',flush=True)
    coefficient,target=read_split(c,'eval')
    cpu_stats={k:v.cpu() for k,v in stats.items()}
    loader=DataLoader(TensorDataset(encode_input(coefficient,cpu_stats),target),
                      batch_size=c['inference']['batch_size'],shuffle=False)
    predicted=[]; relative=[]; elapsed=0
    with torch.no_grad():
        for index,(x,y) in enumerate(loader,start=1):
            x,y=x.to(device),y.to(device)
            if device.type=='cuda': torch.cuda.synchronize()
            tic=time.perf_counter()
            p=decode_output(model(x),stats)
            if device.type=='cuda': torch.cuda.synchronize()
            elapsed+=time.perf_counter()-tic
            r=relative_l2(p,y)
            predicted.append(p.cpu().numpy()); relative.extend(r.cpu().tolist())
            print(f'inference batch={index:02d}/{len(loader)} mean_relative_l2={r.mean().item():.8f}',flush=True)
    prediction=np.concatenate(predicted)
    truth=target.numpy()
    if prediction.shape!=truth.shape or not np.isfinite(prediction).all():
        raise ValueError('Invalid prediction output')
    error=prediction-truth
    output=Path(c['output']['results_dir']);output.mkdir(parents=True,exist_ok=True)
    np.savez_compressed(output/'predictions.npz',coefficient=coefficient.numpy(),target=truth,
                        prediction=prediction,relative_l2=np.asarray(relative),
                        sample_indices=np.arange(len(truth)),checkpoint_epoch=checkpoint['epoch'])
    with (output/'per_sample_metrics.csv').open('w',newline='') as f:
        writer=csv.writer(f);writer.writerow(['sample_index','relative_l2','rmse','mae'])
        for i,r in enumerate(relative):
            writer.writerow([i,r,float(np.sqrt(np.mean(error[i]**2))),float(np.mean(np.abs(error[i])))])
    metrics={'n_eval':len(relative),'resolution':list(truth.shape[1:]),
             'mean_relative_l2':float(np.mean(relative)), 'std_relative_l2':float(np.std(relative)),
             'median_relative_l2':float(np.median(relative)), 'min_relative_l2':float(np.min(relative)),
             'max_relative_l2':float(np.max(relative)), 'rmse':float(np.sqrt(np.mean(error**2))),
             'mae':float(np.mean(np.abs(error))), 'best_epoch':checkpoint['epoch'],
             'checkpoint_eval_relative_l2':checkpoint['eval_relative_l2'],
             'checkpoint':str(checkpoint_path),'inference_seconds':elapsed,
             'milliseconds_per_batch':elapsed/len(loader)*1000,
             'paper_relative_l2':c['reference']['paper_relative_l2'],
             'relative_l2_ratio_to_paper':float(np.mean(relative))/c['reference']['paper_relative_l2'],
             'protocol_note':c['reference']['note']}
    if not np.isclose(metrics['mean_relative_l2'],checkpoint['eval_relative_l2'],rtol=5e-4,atol=1e-7):
        raise ValueError('Loaded best-checkpoint metric does not reproduce saved evaluation')
    atomic_json(output/'metrics.json',metrics)
    print('INFERENCE_COMPLETED '+str(metrics),flush=True)


if __name__=='__main__':
    main()