"""Restore every LOSO fold and emit both methods' 21 complete response maps."""

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.climemu_s2l import FEATURE_COUNT, FORMAT_VERSION, GRID_SHAPE, DualRidge, SharedKernelGPR


def load_checkpoint(path):
    try:
        return torch.load(path, map_location="cpu", weights_only=False)
    except TypeError:
        return torch.load(path, map_location="cpu")


def main():
    config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
    checkpoint = load_checkpoint(ROOT / config["paths"]["checkpoint"])
    if checkpoint["format_version"] != FORMAT_VERSION or tuple(checkpoint["grid_shape"]) != GRID_SHAPE:
        raise ValueError("checkpoint protocol mismatch")
    x = checkpoint["short_response"].double()
    y = checkpoint["long_response"].double()
    required = {"model", "model_config", "format_version"}
    if not required.issubset(checkpoint):
        raise ValueError(f"checkpoint is missing standard fields: {sorted(required - checkpoint.keys())}")
    predictions = {"ridge": [], "gpr": []}
    for expected_fold, state in enumerate(checkpoint["model"]["folds"]):
        if state["fold"] != expected_fold or state["held_out_scenario_id"] != checkpoint["scenario_ids"][expected_fold]:
            raise ValueError("fold/scenario identity mismatch")
        training = state["train_indices"]
        ridge = DualRidge(float(state["ridge_alpha"])).fit(x[training], y[training])
        gpr = SharedKernelGPR(state["gpr_kernel_mode"], config["model"]["gpr"]["jitter"])
        gpr.load_hyperparameters(state["gpr_hyperparameters"]).restore_posterior(x[training], y[training])
        predictions["ridge"].append(ridge.predict(x[expected_fold:expected_fold + 1]).detach().numpy()[0])
        predictions["gpr"].append(gpr.predict(x[expected_fold:expected_fold + 1]).detach().numpy()[0])
    ridge = np.asarray(predictions["ridge"], dtype=np.float32).reshape(21, *GRID_SHAPE)
    gpr = np.asarray(predictions["gpr"], dtype=np.float32).reshape(21, *GRID_SHAPE)
    if ridge.shape != (21, 145, 192) or not np.isfinite(gpr).all():
        raise ValueError("inference did not produce 21 finite full-grid fields")
    source = np.load(ROOT / config["data"]["path"])
    output = ROOT / config["paths"]["inference"]
    output.parent.mkdir(parents=True, exist_ok=True)
    np.savez_compressed(output, format_version=np.array(FORMAT_VERSION),
                        scenario_ids=np.asarray(checkpoint["scenario_ids"]),
                        latitude_deg=source["latitude_deg"], longitude_deg=source["longitude_deg"],
                        short_response=x.numpy().reshape(21, *GRID_SHAPE).astype(np.float32),
                        long_response=y.numpy().reshape(21, *GRID_SHAPE).astype(np.float32),
                        ridge_prediction=ridge, gpr_prediction=gpr)
    print(f"predictions={output.relative_to(ROOT)} methods=2 shape={ridge.shape}")


if __name__ == "__main__":
    main()