from __future__ import annotations

import argparse
import sys
from pathlib import Path

SCRIPT_DIR = Path(__file__).resolve().parent
MODEL_DIR = SCRIPT_DIR.parent / "model"
for module_dir in (SCRIPT_DIR, MODEL_DIR):
    if str(module_dir) not in sys.path:
        sys.path.insert(0, str(module_dir))

import numpy as np
import torch

from common import DEFAULT_CONFIG, active_model_config, load_config, resolve_path
from data import build_loader, load_statistics
from fourcastnet_v2 import (
    FourCastNetV2,
    load_checkpoint,
)


def choose_device() -> torch.device:
    return torch.device("cuda", 0) if torch.cuda.is_available() else torch.device("cpu")


def main() -> None:
    parser = argparse.ArgumentParser(description="Run FourCastNet v2 inference")
    parser.add_argument("--config", default=str(DEFAULT_CONFIG))
    parser.add_argument("--checkpoint")
    args = parser.parse_args()

    config = load_config(args.config)
    inference = config["inference"]

    device = choose_device()
    model = FourCastNetV2(active_model_config(config)).to(device)
    if args.checkpoint:
        checkpoint_path = resolve_path(config, args.checkpoint)
    else:
        checkpoint_path = resolve_path(config, inference["checkpoint_path"])
    result = load_checkpoint(
        model,
        checkpoint_path,
        expected_profile=config["model"]["profile"],
        expected_variables=config["data"]["variables"],
        allowed_stages={"one_step", "finetune"},
        allowed_initializations={"random", "one_step_checkpoint"},
        strict=config["checkpoint"]["strict"],
        map_location=device,
    )
    if result["missing_keys"] or result["unexpected_keys"]:
        print(
            f"missing_keys={result['missing_keys']} "
            f"unexpected_keys={result['unexpected_keys']}"
        )

    loader, _ = build_loader(
        config,
        config["data"]["test_years"],
        train=False,
        distributed=False,
        output_steps=inference["rollout_steps"],
    )
    means, stds = load_statistics(config)
    means = means.numpy()
    stds = stds.numpy()

    output_dir = resolve_path(config, inference["output_dir"])
    output_dir.mkdir(parents=True, exist_ok=True)
    model.eval()
    with torch.no_grad():
        for sample_index, batch in enumerate(loader):
            if sample_index >= inference["max_samples"]:
                break
            inputs = batch[0].to(device)
            targets = batch[1]
            state = inputs
            predictions = []
            for _ in range(inference["rollout_steps"]):
                state = model(state)
                predictions.append(state.cpu())
            prediction = torch.stack(predictions, dim=1).numpy()
            if targets.ndim == 4:
                targets = targets.unsqueeze(1)
            target = targets.numpy()
            input_array = inputs.cpu().numpy()
            if not inference["save_normalized"]:
                prediction = prediction * stds[:, None] + means[:, None]
                target = target * stds[:, None] + means[:, None]
                input_array = input_array * stds + means
            path = output_dir / f"sample_{sample_index:04d}.npz"
            np.savez_compressed(
                path,
                input=input_array,
                prediction=prediction,
                target=target,
                variables=np.asarray(config["data"]["variables"]),
                time_index=np.asarray(batch[4], dtype=str).T,
            )
            print(path)


if __name__ == "__main__":
    main()