"""Reload the selected checkpoint and evaluate all 100 original test cases."""
from __future__ import annotations

import csv
import time

import torch

from common import (DarcyEnergy, arguments, build_model, configuration, error_statistics,
                    evaluate, initialize, output_path, read_data, write_json)


def main():
    parser = arguments(__doc__)
    parser.add_argument("--checkpoint", type=str)
    args = parser.parse_args()
    cfg = configuration(args)
    device = initialize(cfg)
    fields, metadata = read_data(cfg)
    checkpoint_path = output_path(args.checkpoint or cfg["output"]["checkpoint"])
    checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
    if checkpoint["data_sha256"] != metadata["sha256"]:
        raise ValueError("The checkpoint was not trained using this exact dataset")
    if checkpoint["config"]["model"] != cfg["model"]:
        raise ValueError("Model configuration differs from checkpoint")
    model = build_model(cfg, device)
    model.load_state_dict(checkpoint["model_state_dict"], strict=True)
    energy = DarcyEnergy(cfg["data"]["resolution"], cfg["physics"]["source"]).to(device)
    if device.type == "cuda":
        torch.cuda.synchronize()
    started = time.perf_counter()
    prediction, errors, energies = evaluate(model, fields["P_test"].to(device), fields["U_test"].to(device), energy, cfg["training"]["batch_size"])
    if device.type == "cuda":
        torch.cuda.synchronize()
    inference_seconds = time.perf_counter() - started
    boundary = max(prediction[:, 0].abs().max().item(), prediction[:, -1].abs().max().item(),
                   prediction[:, :, 0].abs().max().item(), prediction[:, :, -1].abs().max().item())
    if boundary != 0 or not torch.isfinite(prediction).all():
        raise ValueError("Prediction failed zero-boundary or finiteness check")
    metrics = {**error_statistics(errors), "checkpoint": str(checkpoint_path), "checkpoint_epoch": checkpoint["epoch"],
               "test_energy_mean": energies.double().mean().item(), "boundary_max_abs": boundary,
               "inference_seconds": inference_seconds, "data_sha256": metadata["sha256"],
               "checkpoint_selection": "minimum test-set mean relative L2 across 500 epochs (user requested)",
               "paper_table1_mean_percent": cfg["reference"]["table1_mean_percent"],
               "paper_table1_std_percent": cfg["reference"]["table1_std_percent"]}
    results = output_path(cfg["output"]["results_dir"])
    results.mkdir(parents=True, exist_ok=True)
    torch.save({"raw_permeability": fields["F_test"], "nodal_permeability": fields["P_test"],
                "reference": fields["U_test"], "prediction": prediction, "relative_l2": errors,
                "energy": energies, "checkpoint_epoch": checkpoint["epoch"], "data_sha256": metadata["sha256"]}, results / "test_predictions.pt")
    with (results / "per_sample_metrics.csv").open("w", newline="") as stream:
        writer = csv.writer(stream)
        writer.writerow(["test_index", "relative_l2", "relative_l2_percent", "energy"])
        for index, (error, value) in enumerate(zip(errors.tolist(), energies.tolist())):
            writer.writerow([index, error, 100 * error, value])
            print(f"Test {index:03d}: relative_L2={error:.8f} ({100*error:.4f}%)", flush=True)
    write_json(results / "metrics.json", metrics)
    print(f"Best epoch {checkpoint['epoch']}: test relative L2 = {metrics['relative_l2_mean_percent']:.6f}% "
          f"+/- {metrics['relative_l2_std_percent']:.6f}% (sample std, n=100); boundary={boundary}", flush=True)


if __name__ == "__main__":
    main()