"""Generate a checkpoint-backed CorrDiff ensemble in device-sized batches."""

import argparse
import sys
from pathlib import Path

import numpy as np
import torch
import yaml

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


def scalar(archive, key):
    if key not in archive or archive[key].ndim != 0:
        raise ValueError(f"NPZ metadata {key} must be present as a scalar")
    return str(archive[key].item())


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", default=str(ROOT / "conf/config.yaml"))
    parser.add_argument("--checkpoint")
    args = parser.parse_args()
    config = yaml.safe_load(Path(args.config).read_text(encoding="utf-8"))
    checkpoint_path = Path(args.checkpoint) if args.checkpoint else ROOT / config["paths"]["checkpoint"]
    if not checkpoint_path.is_file():
        raise FileNotFoundError(f"Checkpoint is required: {checkpoint_path}")
    device = torch.device("cuda" if torch.cuda.is_available() and config["runtime"]["device"] != "cpu" else "cpu")
    state = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
    model = CorrDiff(**config["model"]).to(device)
    model.load_state_dict(state["model"])
    model.eval()
    archive = np.load(ROOT / config["data"]["path"])
    protocol, data_source = scalar(archive, "protocol"), scalar(archive, "data_source")
    if protocol != config["data"]["protocol"] or state.get("protocol") != protocol:
        raise ValueError("Data, checkpoint, and configured protocols must match")
    if not data_source or state.get("data_source") != data_source:
        raise ValueError("Data and checkpoint data_source metadata must match")
    coarse = archive[config["data"]["input_key"]]
    target = archive[config["data"]["target_key"]]
    if coarse.ndim != 4 or target.ndim != 4 or tuple(coarse.shape[1:]) != tuple(config["data"]["input_shape"]) or tuple(target.shape[1:]) != tuple(config["data"]["target_shape"]) or len(coarse) != len(target):
        raise ValueError("Invalid CorrDiff NPZ tensor contract")
    member_count = config["sampling"]["ensemble_size"]
    ensemble = np.empty((member_count, len(coarse), *config["data"]["target_shape"]), dtype="float32")
    options = {key: config["sampling"][key] for key in ("steps", "sigma_min", "sigma_max", "rho", "solver")}
    batch_size = config["sampling"]["batch_size"]
    with torch.inference_mode():
        for member in range(member_count):
            torch.manual_seed(config["seed"] + member)
            for start in range(0, len(coarse), batch_size):
                stop = min(start + batch_size, len(coarse))
                batch = torch.from_numpy(coarse[start:stop]).to(device)
                ensemble[member, start:stop] = model.sample(batch, **options).cpu().numpy()
                del batch
    output = ROOT / config["paths"]["predictions"]
    output.parent.mkdir(parents=True, exist_ok=True)
    np.savez_compressed(output, ensemble=ensemble, ensemble_mean=ensemble.mean(0),
                        ensemble_std=ensemble.std(0), target=target, protocol=np.asarray(protocol),
                        data_source=np.asarray(data_source), checkpoint=np.asarray(str(checkpoint_path)))
    print(f"saved={output} ensemble={ensemble.shape} batch_size={batch_size}")


if __name__ == "__main__":
    main()