"""Reload the selected checkpoint and evaluate every short/long test trajectory."""
import json
from pathlib import Path

import torch

from common import (ROOT, Trajectories, build_model, checkpoint_config_matches, device, load_config, output_dir,
                    parser, provenance, rollout_metrics, save_json, setup, sha256)


def main():
    arg = parser(__doc__)
    arg.add_argument("--checkpoint", type=Path)
    args = arg.parse_args()
    cfg = load_config(args.config)
    setup(cfg["seed"])
    target_device = device()
    out = output_dir(cfg)
    checkpoint_path = args.checkpoint or ROOT / cfg["paths"]["best_checkpoint"]
    checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
    status = json.loads((out / "training_status.json").read_text())
    if status["state"] != "completed" or status["completed_epochs"] != 300:
        raise RuntimeError("Final inference requires evidence of completed 300-epoch training")
    if not checkpoint_config_matches(checkpoint["config"], cfg) or checkpoint["epoch"] != status["best_epoch"]:
        raise ValueError("Checkpoint/config/best-epoch provenance mismatch")
    model = build_model(cfg, target_device)
    model.load_state_dict(checkpoint["model_state_dict"], strict=True)
    if not all(torch.isfinite(p).all() for p in model.parameters()):
        raise ValueError("Nonfinite checkpoint parameters")
    information = {"checkpoint": str(checkpoint_path), "sha256": sha256(checkpoint_path),
                   "selected_epoch": checkpoint["epoch"], "completed_epochs": 300,
                   "best_val_mse": checkpoint["best_val_loss"],
                   "parameter_count": sum(p.numel() for p in model.parameters()),
                   "provenance": provenance()}
    print("CHECKPOINT_VERIFIED " + json.dumps(information), flush=True)
    metrics, manifests = {}, {}
    for split in cfg["evaluation"]["splits"]:
        data = Trajectories(cfg, split)
        manifests[split] = data.metadata
        metrics[split] = rollout_metrics(model, data, target_device,
                                         batch_size=cfg["evaluation"]["batch_size"],
                                         save_path=out / split,
                                         example_cases=cfg["evaluation"]["example_cases"])
        print("FINAL_METRICS " + json.dumps(metrics[split]), flush=True)
    save_json(out / "inference_manifest.json", {**information, "data": manifests})
    save_json(out / "metrics.json", metrics)
    print("INFERENCE_COMPLETED all=100 trajectories long_steps=120", flush=True)


if __name__ == "__main__":
    main()