"""Run the four ClimateBench emulators on the held-out scenario."""

import sys
from pathlib import Path

import numpy as np
import torch
import yaml
from torch.utils.data import DataLoader


ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))
from model.climatebench import ClimateBench
from train import ClimateDataset, device_from_config


def main() -> None:
    config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
    device = device_from_config(config)
    checkpoint = torch.load(ROOT / config["paths"]["checkpoint"], map_location=device, weights_only=False)
    if checkpoint["format_version"] != config["data"]["format_version"]:
        raise ValueError("checkpoint and data format versions differ")
    model = ClimateBench(**checkpoint["model_config"]).to(device)
    model.load_state_dict(checkpoint["model"])
    model.eval()
    dataset = ClimateDataset(ROOT / config["data"]["root"] / "test.npz", config)
    loader = DataLoader(dataset, batch_size=1, shuffle=False)
    predictions = []
    with torch.no_grad():
        for inputs, _ in loader:
            predictions.append(model(inputs.to(device)).cpu().numpy())
    prediction = np.concatenate(predictions).astype(np.float32)
    if not np.isfinite(prediction).all():
        raise FloatingPointError("inference produced NaN or Inf")
    source = dataset.data
    output = ROOT / config["paths"]["inference"]
    output.parent.mkdir(parents=True, exist_ok=True)
    np.savez_compressed(output, predictions=prediction, targets=source["targets"], years=source["years"],
                        latitude=source["latitude"], longitude=source["longitude"],
                        target_names=source["target_names"], scenario=source["scenario"],
                        target_aggregation=source["target_aggregation"], storage_layout=np.asarray("NCHW"),
                        format_version=source["format_version"],
                        evaluation_start_year=source["evaluation_start_year"],
                        evaluation_end_year=source["evaluation_end_year"])
    print(f"predictions={output.relative_to(ROOT)} shape={prediction.shape} test_batches={len(predictions)}")


if __name__ == "__main__":
    main()