#!/usr/bin/env python3
import argparse
from pathlib import Path

import numpy as np
import torch

from model.nncam import NNCAM, unscale_output


ROOT = Path(__file__).resolve().parents[1]


def main():
    parser = argparse.ArgumentParser(description="Run offline NNCAM inference.")
    parser.add_argument("--data", type=Path, default=ROOT / "data/nncam_fake.npz")
    parser.add_argument("--checkpoint", type=Path, default=ROOT / "result/checkpoints/nncam.pt")
    parser.add_argument("--output", type=Path, default=ROOT / "result/output/predictions.npz")
    args = parser.parse_args()
    checkpoint = torch.load(args.checkpoint, map_location="cpu", weights_only=True)
    required = {"model", "model_config", "format_version", "normalization"}
    if not required.issubset(checkpoint):
        raise ValueError(f"checkpoint missing {sorted(required - checkpoint.keys())}")
    model = NNCAM(**checkpoint["model_config"])
    model.load_state_dict(checkpoint["model"])
    model.eval()
    with np.load(args.data) as data:
        x, truth, lat, time = (data[name] for name in ("x", "y", "lat", "time"))
    norm = checkpoint["normalization"]
    normalized = (torch.from_numpy(x) - norm["input_mean"]) / norm["input_scale"]
    with torch.no_grad():
        scaled = model(normalized) * norm["target_scale"] + norm["target_mean"]
    prediction = unscale_output(scaled.numpy()).astype(np.float32)
    if prediction.shape != truth.shape or not np.isfinite(prediction).all():
        raise RuntimeError(f"invalid prediction shape or values: {prediction.shape}")
    args.output.parent.mkdir(parents=True, exist_ok=True)
    np.savez_compressed(args.output, input=x, truth=truth, prediction=prediction, lat=lat, time=time)
    print(f"saved {args.output}: prediction={prediction.shape}, finite=true")


if __name__ == "__main__":
    main()