import argparse
from pathlib import Path
import numpy as np
import torch

import sys
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from model.corrdiff import CorrDiff


def main():
    parser = argparse.ArgumentParser(description="Run CorrDiff inference")
    parser.add_argument("--data", default="data/era5_corrdiff.npz")
    parser.add_argument("--checkpoint", default="data/checkpoints/model_bak.pth")
    parser.add_argument("--ensemble-size", type=int, default=1)
    parser.add_argument("--output", default="result/output/predictions.npz")
    args = parser.parse_args()
    data = np.load(args.data)
    coarse = torch.from_numpy(data["input"])
    model = CorrDiff()
    checkpoint = Path(args.checkpoint)
    if checkpoint.exists():
        model.load_state_dict(torch.load(checkpoint, map_location="cpu")["model"])
    model.eval()
    samples = []
    with torch.no_grad():
        for _ in range(args.ensemble_size):
            samples.append(model(coarse).numpy().astype("float32"))
    output = Path(args.output)
    output.parent.mkdir(parents=True, exist_ok=True)
    np.savez_compressed(output, prediction=np.stack(samples), input=data["input"])
    print(f"prediction: {np.stack(samples).shape}")
    print(f"saved: {output}")


if __name__ == "__main__":
    main()