"""Forecast fifteen 6-minute radar frames from five observations."""

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.convlstm import ConvLSTM
from train import RadarDataset, device_from_config


def main():
    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=True)
    model = ConvLSTM(checkpoint["model_config"]).to(device)
    model.load_state_dict(checkpoint["model"])
    model.eval()
    loader = DataLoader(RadarDataset(ROOT / config["data"]["root"] / "test.npz", config), batch_size=1)
    inputs_all, targets_all, predictions_all = [], [], []
    with torch.no_grad():
        for inputs, targets in loader:
            prediction, _ = model(inputs.to(device))
            inputs_all.append(inputs.numpy())
            targets_all.append(targets.numpy())
            predictions_all.append(prediction.cpu().numpy())
    output = ROOT / config["paths"]["inference_dir"] / "predictions.npz"
    output.parent.mkdir(parents=True, exist_ok=True)
    np.savez_compressed(output, inputs=np.concatenate(inputs_all), targets=np.concatenate(targets_all),
                        predictions=np.concatenate(predictions_all),
                        input_lead_minutes=np.arange(-24, 1, int(config["data"]["interval_minutes"]), dtype=np.int64),
                        forecast_lead_minutes=np.arange(1, int(config["data"]["output_frames"]) + 1, dtype=np.int64)
                                              * int(config["data"]["interval_minutes"]),
                        normalized_value_range=np.asarray([0.0, 1.0], np.float32),
                        data_type=np.asarray("normalized_radar_echo_grayscale"))
    print(f"predictions={output.relative_to(ROOT)}")


if __name__ == "__main__":
    main()