"""Reload the best checkpoint and evaluate every held-out validation sample."""

import argparse
import csv
import sys
import time
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
import numpy as np
import torch
from model.lino import LiNO
from common import (compute_device, decode, file_digest, load_checkpoint, load_config,
                    make_inputs, package_path, read_fields, relative_l2, save_json)


@torch.no_grad()
def inference(config, checkpoint=None):
    device = compute_device(config["inference"]["device"])
    checkpoint = package_path(checkpoint or config["training"]["checkpoint"])
    state = load_checkpoint(checkpoint)
    contract = state["data_contract"]
    for key in ["resolution", "interpolation", "input_key", "target_key"]:
        if config["data"][key] != contract[key]:
            raise ValueError(f"Data contract mismatch for {key}")
    data_path = package_path(config["data"]["directory"]) / config["data"]["filename"]
    if file_digest(data_path) != contract["dataset_sha256"]:
        raise ValueError("Dataset content differs from the training data (paths may change)")
    initial, final = read_fields(config)
    ids, stats = state["val_indices"], state["normalizers"]
    model = LiNO(**state["architecture"]).to(device)
    model.load_state_dict(state["model_state"], strict=True)
    model.eval()
    inputs = make_inputs(initial[ids], stats)
    targets = final[ids]
    predictions = []
    started = time.monotonic()
    batch_size = config["inference"]["batch_size"]
    for start in range(0, len(ids), batch_size):
        normalized = model(inputs[start:start+batch_size].to(device))
        predictions.append(decode(normalized, stats).cpu())
        if start % (batch_size * 25) == 0:
            print(f"INFERENCE samples={min(start+batch_size,len(ids))}/{len(ids)} checkpoint_epoch={state['epoch']}", flush=True)
    prediction = torch.cat(predictions)
    errors = relative_l2(prediction, targets)
    mse_by_sample = (prediction-targets).square().flatten(1).mean(1)
    if not torch.isfinite(prediction).all():
        raise FloatingPointError("Nonfinite inference output")
    output = package_path(config["results"]["directory"])
    output.mkdir(parents=True, exist_ok=True)
    np.savez_compressed(output/"predictions.npz", initial=initial[ids].numpy(), target=targets.numpy(),
                        prediction=prediction.numpy(), sample_indices=ids.numpy(), relative_l2=errors.numpy(),
                        x=np.linspace(0,1,config["data"]["resolution"],dtype=np.float32))
    with open(output/"per_sample_metrics.csv", "w", newline="") as stream:
        writer = csv.writer(stream)
        writer.writerow(["validation_index", "source_index", "relative_l2", "mse_physical"])
        writer.writerows((i,int(source),float(rel),float(mse)) for i,(source,rel,mse) in enumerate(zip(ids,errors,mse_by_sample)))
    metric = errors.mean().item()
    metrics = dict(n_validation=len(ids), checkpoint_epoch=state["epoch"], best_epoch=state["best_epoch"],
                   validation_relative_l2=metric, validation_relative_l2_std=errors.std(correction=0).item(),
                   validation_mse_physical=mse_by_sample.mean().item(),
                   validation_rmse_physical=mse_by_sample.mean().sqrt().item(),
                   checkpoint_recorded_relative_l2=state["best_val_relative_l2"],
                   reload_metric_absolute_difference=abs(metric-state["best_val_relative_l2"]),
                   paper_table1_relative_l2=0.0109, ratio_to_paper=metric/0.0109,
                   inference_seconds=time.monotonic()-started,
                   checkpoint_sha256=file_digest(checkpoint), device=torch.cuda.get_device_name(device))
    if metrics["reload_metric_absolute_difference"] > 1e-6:
        raise AssertionError(f"Reloaded metric disagrees with checkpoint: {metrics}")
    save_json(output/"metrics.json", metrics)
    print("BEST CHECKPOINT INFERENCE COMPLETED", metrics, flush=True)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--config", default="config/config.yaml")
    parser.add_argument("--checkpoint", default=None)
    args = parser.parse_args()
    inference(load_config(args.config), args.checkpoint)