"""Compute paper metrics for every scenario and draw maps and error boxes."""

import json
import sys
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]
sys.path.insert(0, str(ROOT))
from model.climemu_s2l import GRID_SHAPE, area_weights, weighted_rmse


REGIONS = {
    "North America": (15, 75, 190, 310), "Europe": (35, 70, 350, 45),
    "South Asia": (5, 35, 60, 100), "East Asia": (20, 55, 100, 150),
    "Arctic": (66, 90, 0, 360), "Northwest Asia": (50, 75, 45, 120),
    "Northern Africa": (15, 35, 340, 55), "Southern Africa": (-35, 0, 10, 55),
    "South America": (-60, 15, 275, 330), "Australia": (-45, -10, 110, 155),
}


def region_mask(latitude, longitude, bounds):
    south, north, west, east = bounds
    latitude_mask = (latitude[:, None] >= south) & (latitude[:, None] <= north)
    longitude_mask = ((longitude[None, :] >= west) & (longitude[None, :] <= east)
                      if west <= east else (longitude[None, :] >= west) | (longitude[None, :] <= east))
    return latitude_mask & longitude_mask


def weighted_mean(field, weights, mask=None):
    selected = np.ones(field.shape, dtype=bool) if mask is None else mask
    local_weights = weights[selected]
    return float(np.sum(field[selected] * local_weights) / np.sum(local_weights))


def main():
    config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
    data = np.load(ROOT / config["paths"]["inference"])
    latitude, longitude = data["latitude_deg"], data["longitude_deg"]
    weights = area_weights(latitude, GRID_SHAPE[1])
    targets = data["long_response"]
    scenario_ids = [str(item) for item in data["scenario_ids"]]
    masks = {name: region_mask(latitude, longitude, bounds) for name, bounds in REGIONS.items()}
    report = {"grid": list(GRID_SHAPE), "scenario_count": 21, "regions": list(REGIONS), "methods": {}}
    box_values = {}
    for method in ("ridge", "gpr"):
        predictions = data[f"{method}_prediction"]
        scenarios = []
        for index, scenario_id in enumerate(scenario_ids):
            target, prediction = targets[index], predictions[index]
            regional = {name: abs(weighted_mean(prediction, weights, mask) - weighted_mean(target, weights, mask))
                        for name, mask in masks.items()}
            scenarios.append({"scenario_id": scenario_id,
                              "area_weighted_grid_rmse": weighted_rmse(target, prediction, weights),
                              "global_mean_absolute_error": abs(weighted_mean(prediction, weights) - weighted_mean(target, weights)),
                              "regional_mean_absolute_error": regional})
        grid_errors = [item["area_weighted_grid_rmse"] for item in scenarios]
        global_errors = [item["global_mean_absolute_error"] for item in scenarios]
        report["methods"][method] = {"scenarios": scenarios,
            "summary": {"mean_grid_rmse": float(np.mean(grid_errors)), "median_grid_rmse": float(np.median(grid_errors)),
                        "mean_global_absolute_error": float(np.mean(global_errors)),
                        "mean_regional_absolute_error": {name: float(np.mean([item["regional_mean_absolute_error"][name]
                                                                             for item in scenarios])) for name in REGIONS}}}
        box_values[method] = [grid_errors, global_errors] + [[item["regional_mean_absolute_error"][name] for item in scenarios]
                                                             for name in REGIONS]
    numeric = [value for method in report["methods"].values() for value in
               (method["summary"]["mean_grid_rmse"], method["summary"]["mean_global_absolute_error"])]
    if not np.isfinite(numeric).all():
        raise FloatingPointError("evaluation metrics are not finite")
    output = ROOT / config["paths"]["evaluation_dir"]
    output.mkdir(parents=True, exist_ok=True)
    (output / "metrics.json").write_text(json.dumps(report, indent=2) + "\n")
    plot_index = scenario_ids.index(config["evaluation"]["plot_scenario_id"])
    fields = [targets[plot_index], data["ridge_prediction"][plot_index], data["gpr_prediction"][plot_index],
              data["ridge_prediction"][plot_index] - targets[plot_index], data["gpr_prediction"][plot_index] - targets[plot_index]]
    titles = ["HadGEM3 target", "Dual Ridge", "Shared-kernel GPR", "Ridge error", "GPR error"]
    figure, axes = plt.subplots(2, 3, figsize=(14, 7), constrained_layout=True)
    for axis, field, title in zip(axes.flat, fields, titles):
        limit = max(abs(np.percentile(field, 1)), abs(np.percentile(field, 99)))
        image = axis.imshow(field, origin="lower", extent=(0, 360, -90, 90), cmap="RdBu_r", vmin=-limit, vmax=limit, aspect="auto")
        axis.set(title=title, xlabel="Longitude", ylabel="Latitude"); figure.colorbar(image, ax=axis, shrink=0.75)
    axes.flat[-1].axis("off")
    figure.suptitle(config["evaluation"]["plot_scenario_id"] + " full 145x192 fields")
    figure.savefig(output / "spatial_fields.png", dpi=150); plt.close(figure)
    labels = ["Grid RMSE", "Global"] + list(REGIONS)
    figure, axes = plt.subplots(1, 2, figsize=(15, 5), sharey=False, constrained_layout=True)
    for axis, method, color in zip(axes, ("ridge", "gpr"), ("#2a6f97", "#c45d35")):
        boxes = axis.boxplot(box_values[method], patch_artist=True, showmeans=True)
        for patch in boxes["boxes"]: patch.set_facecolor(color); patch.set_alpha(0.65)
        axis.set_xticks(range(1, len(labels) + 1), labels, rotation=55, ha="right")
        axis.set(title=method.upper(), ylabel="Absolute error / RMSE (deg C)"); axis.grid(axis="y", alpha=0.25)
    figure.savefig(output / "error_boxplots.png", dpi=150); plt.close(figure)
    print(f"evaluation={output.relative_to(ROOT)} scenarios=21 regions=10 methods=2")


if __name__ == "__main__":
    main()