"""Run bounded ensemble reconstruction from a unified checkpoint."""

from pathlib import Path
import argparse
import json
import sys
import numpy as np
import torch
import yaml

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))

from model.crai_climateextremes import CRAIClimateExtremes


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", type=Path, default=ROOT / "conf/config.yaml")
    args = parser.parse_args()
    config_path = args.config if args.config.is_absolute() else ROOT / args.config
    with open(config_path, encoding="utf-8") as handle:
        cfg = yaml.safe_load(handle)
    data = np.load(ROOT / cfg["data_path"])
    inputs = torch.from_numpy(np.concatenate((data["observed"], data["valid_mask"]), axis=1))
    checkpoint_path = ROOT / cfg["checkpoint_path"]
    if not checkpoint_path.is_file():
        raise FileNotFoundError(f"checkpoint not found: {checkpoint_path}; run scripts/train.py first")
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=True)
    if checkpoint.get("format_version") != "1.0" or not isinstance(checkpoint.get("model"), list):
        raise ValueError(f"unsupported checkpoint format: {checkpoint_path}")
    model_config = checkpoint.get("model_config", {})
    predictions = []
    for state in checkpoint["model"]:
        model = CRAIClimateExtremes(**model_config).to(device)
        model.load_state_dict(state); model.eval()
        with torch.no_grad():
            predictions.append(model(inputs.to(device)).cpu().numpy())
    if not predictions:
        raise ValueError(f"checkpoint contains no ensemble members: {checkpoint_path}")
    members = np.stack(predictions)
    output = ROOT / cfg["output_dir"]
    output.mkdir(parents=True, exist_ok=True)
    np.savez_compressed(
        output / "predictions.npz", prediction=members.mean(0),
        ensemble_std=members.std(0), target=data["target"],
        observed=data["observed"], valid_mask=data["valid_mask"],
        europe_mask=data["europe_mask"], index_ids=data["index_ids"],
        index_names=data["index_names"],
    )
    metadata = {"ensemble_members": len(predictions), "checkpoint_semantics": "member state list in one checkpoint", "output_range": [0, 100]}
    (output / "metadata.json").write_text(json.dumps(metadata, indent=2) + "\n")
    print(f"predicted {inputs.shape[0]} samples with {len(predictions)} members")


if __name__ == "__main__":
    main()