from __future__ import annotations

from typing import Dict, Iterable, List

import numpy as np
import xarray as xr

from src.safe_ops import safe_mean


METRIC_NAMES = ["MAE", "RMSE", "BIAS", "CORR", "RMAE_percent"]


def safe_corr(x: np.ndarray, y: np.ndarray, min_count: int = 3) -> float:
    valid = np.isfinite(x) & np.isfinite(y)
    if valid.sum() < min_count:
        return np.nan
    xv = x[valid]
    yv = y[valid]
    if np.nanstd(xv) == 0 or np.nanstd(yv) == 0:
        return np.nan
    return float(np.corrcoef(xv, yv)[0, 1])


def grid_metrics(pred: np.ndarray, obs: np.ndarray, rmae_percent: bool = True, eps: float = 1e-6) -> Dict[str, np.ndarray]:
    """Calculate metrics at each grid point over time.

    pred, obs: time, lat, lon
    """
    valid = np.isfinite(pred) & np.isfinite(obs)
    err = pred - obs
    abs_err = np.abs(err)
    with np.errstate(invalid="ignore", divide="ignore"):
        n = valid.sum(axis=0)
        mae = np.where(n > 0, np.nansum(np.where(valid, abs_err, np.nan), axis=0) / n, np.nan)
        rmse = np.where(n > 0, np.sqrt(np.nansum(np.where(valid, err * err, np.nan), axis=0) / n), np.nan)
        bias = np.where(n > 0, np.nansum(np.where(valid, err, np.nan), axis=0) / n, np.nan)
        denom = np.nansum(np.where(valid, obs, np.nan), axis=0)
        rmae = np.where(np.abs(denom) > eps, np.nansum(np.where(valid, abs_err, np.nan), axis=0) / denom, np.nan)
        if rmae_percent:
            rmae = 100.0 * rmae

    nt, nlat, nlon = pred.shape
    corr = np.full((nlat, nlon), np.nan, dtype=np.float32)
    for j in range(nlat):
        for i in range(nlon):
            corr[j, i] = safe_corr(pred[:, j, i], obs[:, j, i])

    return {
        "MAE": mae.astype(np.float32),
        "RMSE": rmse.astype(np.float32),
        "BIAS": bias.astype(np.float32),
        "CORR": corr.astype(np.float32),
        "RMAE_percent": rmae.astype(np.float32),
    }


def metrics_to_dataset(metrics: Dict[str, np.ndarray], lat: np.ndarray, lon: np.ndarray) -> xr.Dataset:
    ds = xr.Dataset(coords={"lat": lat, "lon": lon})
    for name, arr in metrics.items():
        ds[name] = (("lat", "lon"), arr.astype(np.float32))
    ds["RMAE_percent"].attrs["long_name"] = "Relative mean absolute error"
    ds["RMAE_percent"].attrs["units"] = "%"
    ds["CORR"].attrs["long_name"] = "Temporal correlation coefficient"
    ds["MAE"].attrs["long_name"] = "Mean absolute error"
    ds["RMSE"].attrs["long_name"] = "Root mean square error"
    ds["BIAS"].attrs["long_name"] = "Mean error, prediction minus observation"
    return ds


def summarize_by_region(
    metrics: Dict[str, np.ndarray],
    reg: np.ndarray,
    regions: Iterable[int],
    meta: Dict[str, object],
) -> List[Dict[str, object]]:
    rows: List[Dict[str, object]] = []
    target_name = str(meta.get("target", ""))
    for region in regions:
        rmask = np.isfinite(reg) & (reg == float(region))
        for metric_name, arr in metrics.items():
            if target_name != "pr_o" and metric_name == "RMAE_percent":
                continue
            vals = arr[rmask]
            if np.isfinite(vals).any():
                value = float(safe_mean(vals, min_count=1, dtype=float))
                median = float(np.nanmedian(vals))
            else:
                value = np.nan
                median = np.nan
            row = dict(meta)
            row.update({
                "region": f"R{int(region)}",
                "metric": metric_name,
                "mean_value": value,
                "median_value": median,
                "n_grid": int(np.isfinite(vals).sum()),
            })
            rows.append(row)

    # Vietnam-wide summary across all valid regions.
    vmask = np.isfinite(reg) & (reg > 0)
    for metric_name, arr in metrics.items():
        if target_name != "pr_o" and metric_name == "RMAE_percent":
            continue
        vals = arr[vmask]
        row = dict(meta)
        row.update({
            "region": "ALL",
            "metric": metric_name,
            "mean_value": float(safe_mean(vals, min_count=1, dtype=float)) if np.isfinite(vals).any() else np.nan,
            "median_value": float(np.nanmedian(vals)) if np.isfinite(vals).any() else np.nan,
            "n_grid": int(np.isfinite(vals).sum()),
        })
        rows.append(row)
    return rows