"""Load validation-selected weights and evaluate the last 200 independent cases."""
import csv
import os
import time

import numpy as np
from common import (parser, load_config, seed_everything, device_for, resolve,
                    load_splits, load_best, evaluate, save_json, digest, update_run_manifest)


def main():
    args = parser(__doc__).parse_args()
    cfg = load_config(args.config)
    seed_everything(cfg['seed'], cfg['training']['num_threads'])
    device = device_for(cfg['inference']['device'])
    model, checkpoint, checkpoint_path = load_best(cfg, device)
    splits, times, data_path = load_splits(cfg, ['test'])
    if checkpoint['data_sha256'] != digest(data_path):
        raise ValueError('Raw test data checksum differs from the training source')
    u, y, ids = splits['test']
    start = time.monotonic()
    summary, pred, rmse_i, relative_i, latent = evaluate(
        model, u.to(device), y.to(device), times, cfg['inference']['batch_size'], True)
    summary.update(split='independent_test', indices=[int(ids[0]), int(ids[-1])],
                   checkpoint_epoch=checkpoint['epoch'], checkpoint_sha256=digest(checkpoint_path),
                   checkpoint_validation_mse=checkpoint['validation_mse'],
                   selection_split=checkpoint['selection_split'],
                   paper_rmse_mean=0.00134, paper_rmse_std=0.00162,
                   elapsed_seconds=time.monotonic() - start,
                   slurm_job_id=os.environ.get('SLURM_JOB_ID'))
    result_dir = resolve(cfg['paths']['results'])
    result_dir.mkdir(parents=True, exist_ok=True)
    np.savez_compressed(result_dir / 'test_predictions.npz', prediction=pred.numpy(),
                        reference=y.numpy(), initial=u.numpy(), latent=latent.numpy(),
                        sample_ids=ids, x=np.linspace(0, 1, 33), times=times.numpy())
    with (result_dir / 'test_per_sample.csv').open('w', newline='') as file:
        writer = csv.writer(file)
        writer.writerow(['sample_id', 'rmse', 'relative_l2'])
        writer.writerows(zip(ids, rmse_i, relative_i))
    save_json(result_dir / 'test_metrics.json', summary)
    update_run_manifest(cfg, status='inference_completed', test_samples=len(ids),
                        test_indices=[int(ids[0]), int(ids[-1])],
                        test_metrics_file='test_metrics.json')
    print(f'INDEPENDENT TEST n={len(ids)} RMSE mean={summary["trajectory_rmse_mean"]:.9e} '
          f'std(ddof1)={summary["trajectory_rmse_std_ddof1"]:.9e} '
          f'global_RMSE={summary["global_rmse"]:.9e} '
          f'relative_L2={summary["relative_l2_mean"]:.9e}', flush=True)
    print(f'Table 1(a) units 1e-3: {summary["trajectory_rmse_mean"] * 1000:.4f} '
          f'({summary["trajectory_rmse_std_ddof1"] * 1000:.4f}); paper 1.34 (1.62)', flush=True)


if __name__ == '__main__':
    main()