"""Checkpoint-backed one-step and autoregressive RainNet inference."""

import json
from pathlib import Path
import sys

import h5py
import numpy as np
import torch
import torch.nn.functional as F
import yaml

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))
from model.rainnet import build_rainnet


def load_config():
    with (ROOT / "conf/config.yaml").open(encoding="utf-8") as handle:
        return yaml.safe_load(handle)


def inverse_transform(tensor):
    return torch.clamp(torch.exp(tensor) - 0.01, min=0.0)


def describe(name, array):
    print(
        f"{name}: shape={array.shape}, dtype={array.dtype}, "
        f"min={array.min():.8f}, max={array.max():.8f}, mean={array.mean():.8f}"
    )


def main():
    config = load_config()
    torch.manual_seed(config["seed"])
    requested_device = config["device"]
    if requested_device == "auto":
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    else:
        device = torch.device(requested_device)
    checkpoint_path = ROOT / config["inference"]["checkpoint"]
    if not checkpoint_path.exists():
        raise FileNotFoundError(f"Required checkpoint does not exist: {checkpoint_path}")
    model = build_rainnet(**config["model"]).to(device)
    checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False)
    model.load_state_dict(checkpoint["model_state_dict"])
    model.eval()
    print("checkpoint_loaded: True")
    print(f"checkpoint path: {checkpoint_path}")

    data = config["data"]
    data_path = ROOT / data["path"]
    with h5py.File(data_path, "r") as handle:
        keys = sorted(handle.keys())
        test_start = data["train_frames"] + data["val_frames"]
        test_keys = keys[test_start : test_start + data["test_frames"]]
        raw = np.stack([handle[key][...] for key in test_keys]).astype(np.float32)
    needed = data["input_steps"] + config["inference"]["rollout_steps"]
    if len(raw) < needed:
        raise ValueError(f"Inference needs {needed} test frames, found {len(raw)}")

    inputs = torch.from_numpy(raw[: data["input_steps"]]).unsqueeze(0).to(device)
    targets = raw[data["input_steps"] : needed]
    log_window = torch.log(inputs + 0.01)
    pad_h = data["padded_height"] - data["raw_height"]
    pad_w = data["padded_width"] - data["raw_width"]
    pad = (pad_w // 2, pad_w - pad_w // 2, pad_h // 2, pad_h - pad_h // 2)
    crop = (pad_h // 2, pad_w // 2)
    predictions = []
    with torch.inference_mode():
        for _ in range(config["inference"]["rollout_steps"]):
            padded = F.pad(log_window, pad, mode="reflect")
            padded_prediction = model(padded)
            prediction = padded_prediction[
                :, :, crop[0] : crop[0] + data["raw_height"], crop[1] : crop[1] + data["raw_width"]
            ]
            predictions.append(inverse_transform(prediction).squeeze(0).squeeze(0).cpu().numpy())
            log_window = torch.cat((log_window[:, 1:], prediction), dim=1)
    predictions = np.stack(predictions).astype(np.float32)
    persistence = np.repeat(raw[data["input_steps"] - 1][None], len(predictions), axis=0).astype(np.float32)
    output_dir = ROOT / config["inference"]["output_dir"]
    output_dir.mkdir(parents=True, exist_ok=True)
    np.save(output_dir / "inputs.npy", raw[: data["input_steps"]])
    np.save(output_dir / "predictions.npy", predictions)
    np.save(output_dir / "targets.npy", targets)
    np.save(output_dir / "persistence.npy", persistence)
    metadata = {
        "units": "mm/5min",
        "interval_minutes": data["interval_minutes"],
        "rollout_steps": len(predictions),
        "input_keys": test_keys[: data["input_steps"]],
        "target_keys": test_keys[data["input_steps"] : needed],
        "checkpoint": str(checkpoint_path),
    }
    with (output_dir / "metadata.json").open("w", encoding="utf-8") as handle:
        json.dump(metadata, handle, indent=2)
    describe("inputs", raw[: data["input_steps"]])
    describe("predictions", predictions)
    describe("targets", targets)
    describe("persistence", persistence)
    print(f"One-step inference shape: {predictions[:1].shape}")
    print(f"Autoregressive rollout steps: {len(predictions)}")


if __name__ == "__main__":
    main()