"""Generate complete test probabilities, classes, centroids, and years."""

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.cesm_seasonal_ml import (CHECKPOINT_FORMAT_VERSION, DATA_FORMAT_VERSION, FeedForwardNN,
                                    SeasonalLSTM, StablePrecipKMeans)


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"])
    data = np.load(ROOT / config["data"]["path"])
    required = {"format_version", "model", "model_config", "data_spec", "season"}
    if not required.issubset(checkpoint):
        raise ValueError(f"checkpoint missing required top-level keys: {sorted(required - checkpoint.keys())}")
    if checkpoint["format_version"] != CHECKPOINT_FORMAT_VERSION:
        raise ValueError("checkpoint format version mismatch")
    spec, model_config = checkpoint["data_spec"], checkpoint["model_config"]
    if str(data["format_version"]) != DATA_FORMAT_VERSION or spec["format_version"] != DATA_FORMAT_VERSION:
        raise ValueError("data format version mismatch")
    season = checkpoint["season"]
    if season != spec["season"] or season not in [str(value) for value in data["seasons"]]:
        raise ValueError("checkpoint/data season mismatch")
    checks = (("rf_manifest", spec["rf_manifest"]), ("nn_manifest", spec["nn_manifest"]),
              ("class_names", spec["class_names"]))
    for key, expected in checks:
        if data[key].tolist() != expected:
            raise ValueError(f"checkpoint/data {key} mismatch")
    shapes = (("rf_features", "rf_shape"), ("nn_features", "nn_shape"),
              ("eof_sequence", "eof_shape"), ("precipitation", "precipitation_shape"))
    for key, shape_key in shapes:
        if list(data[key].shape) != spec[shape_key]:
            raise ValueError(f"checkpoint/data {key} shape mismatch")
    season_index = [str(x) for x in data["seasons"]].index(season)
    test = data["split"] == 2
    cluster = StablePrecipKMeans.from_state_dict(checkpoint["cluster"])
    target = cluster.transform(data["precipitation"][season_index, test])
    probabilities = {}
    models, settings = checkpoint["model"], model_config["settings"]
    if "rf" in models:
        probabilities["rf"] = models["rf"].predict_proba(data["rf_features"][season_index, test])
    if "xgboost" in models:
        probabilities["xgboost"] = models["xgboost"].predict_proba(data["rf_features"][season_index, test])
    if "nn" in models:
        model = FeedForwardNN(input_size=model_config["nn_features"], hidden=tuple(settings["nn"]["hidden"]), dropout=settings["nn"]["dropout"], classes=model_config["classes"])
        model.load_state_dict(models["nn"]); model.eval()
        with torch.no_grad(): probabilities["nn"] = torch.softmax(model(torch.from_numpy(data["nn_features"][season_index, test]).float()), 1).numpy()
    if "lstm" in models:
        lstm = settings["lstm"]
        model = SeasonalLSTM(input_size=model_config["eof_channels"], hidden_size=lstm["hidden_size"], dense_size=lstm["dense_size"], dropout=lstm["dropout"], classes=model_config["classes"])
        model.load_state_dict(models["lstm"]); model.eval()
        sequence = data["eof_sequence"][season_index, test, -model_config["history_length"]:]
        with torch.no_grad(): probabilities["lstm"] = torch.softmax(model(torch.from_numpy(sequence).float()), 1).numpy()
    primary = model_config["primary"] if model_config["primary"] in probabilities else next(iter(probabilities))
    output = ROOT / config["paths"]["inference"]
    output.parent.mkdir(parents=True, exist_ok=True)
    payload = {"format_version": DATA_FORMAT_VERSION, "season": season, "years": data["years"][test], "target": target, "centroids": cluster.centroids_,
               "latitude": data["latitude"], "longitude": data["longitude"], "primary_model": primary,
               "class_names": data["class_names"], "target_precipitation": data["precipitation"][season_index, test]}
    for name, probability in probabilities.items():
        payload[f"probability_{name}"] = probability
        payload[f"prediction_{name}"] = probability.argmax(axis=1)
    np.savez_compressed(output, **payload)
    print(f"predictions={output.relative_to(ROOT)} season={season} samples={test.sum()} primary={primary}")


if __name__ == "__main__":
    main()