from __future__ import annotations

from pathlib import Path
from typing import Dict, Iterable, List

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


def plot_metric_map(
    lon: np.ndarray,
    lat: np.ndarray,
    arr: np.ndarray,
    title: str,
    output_file: str | Path,
    units: str = "",
) -> None:
    output_file = Path(output_file)
    output_file.parent.mkdir(parents=True, exist_ok=True)
    fig, ax = plt.subplots(figsize=(6.5, 8.0), constrained_layout=True)
    mesh = ax.pcolormesh(lon, lat, arr, shading="auto")
    cb = fig.colorbar(mesh, ax=ax, shrink=0.82)
    if units:
        cb.set_label(units)
    ax.set_xlabel("Longitude")
    ax.set_ylabel("Latitude")
    ax.set_title(title)
    ax.set_xlim(float(np.nanmin(lon)), float(np.nanmax(lon)))
    ax.set_ylim(float(np.nanmin(lat)), float(np.nanmax(lat)))
    fig.savefig(output_file, dpi=150)
    plt.close(fig)


def plot_all_metric_maps(
    lon: np.ndarray,
    lat: np.ndarray,
    metrics: Dict[str, np.ndarray],
    meta_label: str,
    output_dir: str | Path,
) -> List[Path]:
    out_files: List[Path] = []
    units = {"RMAE_percent": "%", "CORR": "", "MAE": "", "RMSE": "", "BIAS": ""}
    for metric, arr in metrics.items():
        fname = Path(output_dir) / f"map_{meta_label}_{metric}.png"
        plot_metric_map(lon, lat, arr, f"{meta_label} {metric}", fname, units.get(metric, ""))
        out_files.append(fname)
    return out_files


def plot_summary_bars(summary_csv: str | Path, output_file: str | Path, target: str, primary_metric: str) -> None:
    df = pd.read_csv(summary_csv)
    df = df[(df["target"] == target) & (df["metric"] == primary_metric) & (df["region"] == "ALL")]
    if df.empty:
        return
    # Aggregate over source/lead if many rows exist.
    grouped = df.groupby(["method", "spatial_mode", "predictor_set"], as_index=False)["mean_value"].mean()
    grouped["label"] = grouped["method"] + "|" + grouped["spatial_mode"] + "|" + grouped["predictor_set"]
    grouped = grouped.sort_values("mean_value")
    fig, ax = plt.subplots(figsize=(10, max(4, 0.25 * len(grouped))), constrained_layout=True)
    ax.barh(grouped["label"], grouped["mean_value"])
    ax.set_xlabel(primary_metric)
    ax.set_title(f"{target}: mean {primary_metric} over ALL regions")
    fig.savefig(output_file, dpi=150)
    plt.close(fig)