"""Run ensemble inference for all held-out synthetic daily maps."""

import sys
from pathlib import Path

import numpy as np
import torch


ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))
from model.precipdd import ensemble_predict, load_config, load_ensemble, validate_archive


def main():
    config = load_config(ROOT / "conf/config.yaml")
    device = torch.device("cuda" if torch.cuda.is_available() and config["runtime"]["device"] != "cpu" else "cpu")
    data = np.load(ROOT / config["paths"]["data"])
    validate_archive(data)
    models, checkpoint = load_ensemble(ROOT / config["paths"]["checkpoint"], device)
    mask = data["split"] == 2
    fields = torch.from_numpy(data["precipitation"][mask]).float().to(device)
    prediction = ensemble_predict(models, fields, config["training"]["batch_size"]).cpu().numpy()
    output = ROOT / config["paths"]["predictions"]
    output.parent.mkdir(parents=True, exist_ok=True)
    np.savez_compressed(output, format_version=data["format_version"], prediction=prediction, target=data["agmt"][mask],
                        precipitation=data["precipitation"][mask], year=data["year"][mask], day_of_year=data["day_of_year"][mask],
                        latitude=data["latitude"], longitude=data["longitude"], ensemble_members=np.array(len(models)),
                        checkpoint_world_size=np.array(checkpoint["world_size"]))
    print(f"predictions={output.relative_to(ROOT)} days={len(prediction)} members={len(models)}")


if __name__ == "__main__":
    main()