"""Run complete-grid inference and preserve every sample in one NPZ."""

import sys
from pathlib import Path

import numpy as np
import torch
import yaml


ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))
from model.precipitationsrcnn import FORMAT_VERSION, MODEL_NAME, bilinear_input, build_model


def load_checkpoint(path):
    try:
        return torch.load(path, map_location="cpu", weights_only=True)
    except TypeError:
        return torch.load(path, map_location="cpu")


def main():
    config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
    data = np.load(ROOT / config["data"]["root"] / "daily_precipitation.npz")
    checkpoint = load_checkpoint(ROOT / config["paths"]["checkpoint"])
    if checkpoint["format_version"] != FORMAT_VERSION or checkpoint["model_name"] != MODEL_NAME:
        raise ValueError("checkpoint version/model mismatch")
    if checkpoint["target_grid"] != [216, 488] or checkpoint["model_config"] != config["model"]:
        raise ValueError("checkpoint shape/config mismatch")
    model = build_model(checkpoint["model_config"])
    model.load_state_dict(checkpoint["model"]); model.eval()
    coarse = torch.from_numpy(data["coarse_precipitation"])
    elevation = torch.from_numpy(np.repeat(data["elevation"], len(coarse), axis=0))
    inputs = bilinear_input(coarse, elevation, (216, 488))
    predictions = []
    with torch.no_grad():
        for index in range(len(inputs)):
            predictions.append(model(inputs[index:index + 1]).numpy())
    prediction = np.concatenate(predictions).astype(np.float32)
    if prediction.shape != data["target_precipitation"].shape or not np.isfinite(prediction).all():
        raise ValueError("incomplete or invalid inference output")
    output = ROOT / config["paths"]["inference"]
    output.parent.mkdir(parents=True, exist_ok=True)
    np.savez_compressed(output, format_version=np.array(FORMAT_VERSION),
                        prediction=prediction, target=data["target_precipitation"],
                        bilinear_precipitation=inputs[:, :1].numpy(), elevation=data["elevation"],
                        coarse_precipitation=data["coarse_precipitation"],
                        timestamps=data["timestamps"], years=data["years"], units=data["units"])
    print(f"predictions={output.relative_to(ROOT)} shape={prediction.shape}")


if __name__ == "__main__":
    main()