"""Compute ClimateBench spatial/global NRMSE and plot all four targets."""

import json
from pathlib import Path

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import yaml


ROOT = Path(__file__).resolve().parents[1]


def area_mean(field: np.ndarray, latitude: np.ndarray) -> np.ndarray:
    weights = np.cos(np.deg2rad(latitude)).clip(0)[:, None]
    return np.sum(field * weights, axis=(-2, -1)) / (weights.sum() * field.shape[-1])


def climatebench_metrics(prediction: np.ndarray, target: np.ndarray, latitude: np.ndarray, alpha: float) -> dict:
    normalization = float(area_mean(target, latitude).mean())
    if abs(normalization) <= 1e-12:
        raise ValueError("area-weighted target mean must be nonzero for NRMSE normalization")
    temporal_rmse = np.sqrt(np.mean((prediction - target) ** 2, axis=0))
    spatial = float(area_mean(temporal_rmse, latitude) / normalization)
    predicted_global = area_mean(prediction, latitude)
    target_global = area_mean(target, latitude)
    global_nrmse = float(np.sqrt(np.mean((predicted_global - target_global) ** 2)) / normalization)
    return {"spatial_nrmse": spatial, "global_nrmse": global_nrmse,
            "total_nrmse": spatial + alpha * global_nrmse}


def main() -> None:
    config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
    data = np.load(ROOT / config["paths"]["inference"])
    if str(data["format_version"]) != config["data"]["format_version"]:
        raise ValueError("incompatible prediction format")
    start, end = int(data["evaluation_start_year"]), int(data["evaluation_end_year"])
    if (start, end) != (int(config["evaluation"]["start_year"]), int(config["evaluation"]["end_year"])):
        raise ValueError("evaluation period metadata must be 2080-2100")
    names = data["target_names"].tolist()
    alpha = float(config["evaluation"]["global_weight"])
    metrics = {name: climatebench_metrics(data["predictions"][:, index], data["targets"][:, index],
                                           data["latitude"], alpha) for index, name in enumerate(names)}
    report = {"variables": metrics, "mean_total_nrmse": float(np.mean([value["total_nrmse"] for value in metrics.values()])),
              "evaluation_protocol": {"scenario": str(data["scenario"]), "period": f"{start}-{end}",
                                      "target_aggregation": str(data["target_aggregation"]),
                                      "spatial_metric": "cosine-area mean of grid-cell temporal RMSE",
                                      "global_weighting": "cosine latitude cell area",
                                      "total": "spatial_nrmse + 5 * global_nrmse",
                                      "normalization": "cosine-area and sample mean of target field"}}
    numeric = [number for values in metrics.values() for number in values.values()]
    if not np.isfinite(numeric).all():
        raise FloatingPointError("evaluation contains NaN or Inf")
    output = ROOT / config["paths"]["evaluation_dir"]
    output.mkdir(parents=True, exist_ok=True)
    (output / "metrics.json").write_text(json.dumps(report, indent=2) + "\n")
    figure, axes = plt.subplots(4, 3, figsize=(12, 12), constrained_layout=True)
    longitude, latitude = data["longitude"], data["latitude"]
    for index, name in enumerate(names):
        fields = (data["targets"][:, index].mean(0), data["predictions"][:, index].mean(0),
                  (data["predictions"][:, index] - data["targets"][:, index]).mean(0))
        for column, (field, title) in enumerate(zip(fields, ("Target", "Prediction", "Error"))):
            image = axes[index, column].pcolormesh(longitude, latitude, field, shading="auto", cmap="coolwarm")
            axes[index, column].set(title=f"{name}: {title}", xlabel="Longitude", ylabel="Latitude")
            figure.colorbar(image, ax=axes[index, column], shrink=0.75)
    figure.suptitle(f"ClimateBench {str(data['scenario'])} {start}-{end}")
    figure.savefig(output / "four_targets.png", dpi=150)
    plt.close(figure)
    print(f"evaluation={output.relative_to(ROOT)} variables={','.join(names)}")


if __name__ == "__main__":
    main()