"""Compute per-lead latitude-weighted model and persistence RMSE/MBE."""

import json
from pathlib import Path
import sys

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.fuxi_ocean import ocean_channel_depth_indices, weighted_error_sums


def aggregate(prediction, truth, latitude, depth_mask):
    channel_depth = ocean_channel_depth_indices()
    mask = np.take(depth_mask, channel_depth, axis=1)
    squared, bias, weight = weighted_error_sums(prediction, truth, latitude, mask)
    return np.sqrt(squared / np.maximum(weight, 1)), bias / np.maximum(weight, 1)


def main():
    config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
    data = np.load(ROOT / config["paths"]["inference"])
    model_rmse, model_mbe, persistence_rmse, persistence_mbe = [], [], [], []
    for lead in range(data["prediction"].shape[1]):
        rmse, mbe = aggregate(data["prediction"][:, lead], data["truth"][:, lead], data["latitude_deg"], data["depth_mask"])
        prmse, pmbe = aggregate(data["initial"], data["truth"][:, lead], data["latitude_deg"], data["depth_mask"])
        model_rmse.append(rmse); model_mbe.append(mbe); persistence_rmse.append(prmse); persistence_mbe.append(pmbe)
    arrays = [np.asarray(value) for value in (model_rmse, model_mbe, persistence_rmse, persistence_mbe)]
    lead_hours = data["lead_hours"].tolist()
    groups = json.loads(str(data["variable_groups"]))
    def grouped(values, i):
        return {name: {"unit": spec["unit"], "rmse": float(np.mean(values[0][i, slice(*spec["channels"])])),
                       "mbe": float(np.mean(values[1][i, slice(*spec["channels"])]))} for name, spec in groups.items()}
    metrics = {"output_kind": str(data["output_kind"]), "format_version": str(data["format_version"]),
               "checkpoint_source": str(data["checkpoint_source"]), "sample_count": int(data["sample_count"]),
               "coverage_fraction": float(data["coverage_fraction"]), "is_complete_global": bool(data["is_complete_global"]),
               "synthetic": bool(data["synthetic"]), "aggregation": "latitude-weighted sampled tiles; never across unit groups",
               "output_shape": data["output_shape"].tolist(), "variable_groups": groups,
               "per_lead": [{"lead_hours": hour, "model_by_group": grouped(arrays[:2], i),
                              "persistence_by_group": grouped(arrays[2:], i),
                              "model_by_channel": {"rmse": arrays[0][i].tolist(), "mbe": arrays[1][i].tolist()},
                              "persistence_by_channel": {"rmse": arrays[2][i].tolist(), "mbe": arrays[3][i].tolist()}}
                             for i, hour in enumerate(lead_hours)]}
    path = ROOT / config["paths"]["evaluation_metrics"]; path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(metrics, indent=2) + "\n")
    figure, axes = plt.subplots(1, 5, figsize=(18, 3.5))
    for axis, (name, spec) in zip(axes, groups.items()):
        channel_slice = slice(*spec["channels"])
        axis.plot(lead_hours, arrays[0][:, channel_slice].mean(1), "o-", label="model")
        axis.plot(lead_hours, arrays[2][:, channel_slice].mean(1), "s--", label="persistence")
        axis.set(title=name, xlabel="Lead (h)", ylabel=f"RMSE ({spec['unit']})"); axis.legend()
    figure.tight_layout()
    plot = ROOT / config["paths"]["evaluation_plot"]; figure.savefig(plot, dpi=160); plt.close(figure)
    print(f"metrics={path.relative_to(ROOT)} plot={plot.relative_to(ROOT)} leads={lead_hours}")


if __name__ == "__main__":
    main()