from __future__ import annotations

from typing import Sequence

import numpy as np
import xarray as xr


def compute_mode_correlation_map(field: xr.DataArray, scores: xr.DataArray, mode: int) -> xr.DataArray:
    s = scores.sel(mode=mode)
    s, field = xr.align(s, field, join="inner")
    return xr.corr(field, s, dim="time").rename(f"mode{mode}_corr")


# =========================================================
# Constants
# =========================================================
DEFAULT_CATEGORY_LABELS = ("below", "near", "above")
DEFAULT_CATEGORY_CODES = (0, 1, 2)


# =========================================================
# Helpers
# =========================================================
def _require_dims(da: xr.DataArray, dims: Sequence[str], name: str) -> None:
    missing = [d for d in dims if d not in da.dims]
    if missing:
        raise ValueError(f"{name} is missing required dims: {missing}. Found dims={da.dims}")


def _validate_probabilities(
    prob_fcst: xr.DataArray,
    *,
    category_dim: str,
    sample_dim: str,
) -> None:
    _require_dims(prob_fcst, (category_dim, sample_dim, "lat", "lon"), "prob_fcst")

    ncat = prob_fcst.sizes[category_dim]
    if ncat != 3:
        raise ValueError(
            f"prob_fcst must have exactly 3 categories along {category_dim!r}, got {ncat}"
        )


def _safe_nanquantile(da: xr.DataArray, q: float, dim: str) -> xr.DataArray:
    """
    Quantile helper that does not crash on fully-NaN slices.
    """
    if not 0.0 <= q <= 1.0:
        raise ValueError(f"q must be between 0 and 1, got {q}")

    _require_dims(da, (dim,), "da")

    return xr.apply_ufunc(
        lambda x: np.nanquantile(x, q, axis=-1),
        da.transpose(..., dim),
        input_core_dims=[[dim]],
        output_core_dims=[[]],
        vectorize=True,
        dask="parallelized",
        output_dtypes=[float],
    )


def _get_category_labels(prob_fcst: xr.DataArray, category_dim: str) -> list[str]:
    labels = [str(v) for v in prob_fcst[category_dim].values]
    if len(labels) != 3:
        raise ValueError(f"Expected 3 category labels, got {labels}")
    return labels


def _resolve_category_label(
    prob_fcst: xr.DataArray,
    requested: str,
    category_dim: str,
) -> str:
    labels = _get_category_labels(prob_fcst, category_dim)
    if requested in labels:
        return requested

    raise ValueError(
        f"Category label {requested!r} not found in prob_fcst[{category_dim!r}]. "
        f"Available labels: {labels}"
    )


def _count_valid_pairs(a: xr.DataArray, b: xr.DataArray, sample_dim: str) -> xr.DataArray:
    valid = xr.apply_ufunc(np.isfinite, a) & xr.apply_ufunc(np.isfinite, b)
    return valid.sum(dim=sample_dim)


# =========================================================
# Masks
# =========================================================
def build_metric_mask(
    y_true: xr.DataArray,
    y_pred: xr.DataArray,
    *,
    sample_dim: str = "time",
    min_valid_time: int = 3,
) -> xr.DataArray:
    """
    Keep only cells with enough valid paired samples across time.
    """
    if min_valid_time < 1:
        raise ValueError("min_valid_time must be >= 1")

    _require_dims(y_true, (sample_dim, "lat", "lon"), "y_true")
    _require_dims(y_pred, (sample_dim, "lat", "lon"), "y_pred")

    valid_count = _count_valid_pairs(y_true, y_pred, sample_dim=sample_dim)

    return xr.DataArray(
        (valid_count >= min_valid_time).values,
        dims=("lat", "lon"),
        coords={"lat": y_true["lat"], "lon": y_true["lon"]},
        name="metric_mask",
    )


def build_prob_metric_mask(
    obs_class: xr.DataArray,
    prob_fcst: xr.DataArray,
    *,
    sample_dim: str = "time",
    category_dim: str = "category",
    min_valid_time: int = 3,
    require_sum_to_one: bool = False,
    sum_tolerance: float = 1e-6,
) -> xr.DataArray:
    """
    Keep only cells with enough valid probability-observation pairs.

    A time/lat/lon pair is valid when:
      - obs_class is finite
      - all category probabilities are finite
      - optionally, category probabilities sum to 1 within tolerance
    """
    if min_valid_time < 1:
        raise ValueError("min_valid_time must be >= 1")
    if sum_tolerance < 0:
        raise ValueError("sum_tolerance must be >= 0")

    _require_dims(obs_class, (sample_dim, "lat", "lon"), "obs_class")
    _validate_probabilities(prob_fcst, category_dim=category_dim, sample_dim=sample_dim)

    prob_valid = xr.apply_ufunc(np.isfinite, prob_fcst).all(dim=category_dim)
    obs_valid = xr.apply_ufunc(np.isfinite, obs_class)

    valid_pairs = prob_valid & obs_valid

    if require_sum_to_one:
        prob_sum = prob_fcst.sum(dim=category_dim)
        sum_ok = np.abs(prob_sum - 1.0) <= sum_tolerance
        valid_pairs = valid_pairs & sum_ok

    valid_count = valid_pairs.sum(dim=sample_dim)

    return xr.DataArray(
        (valid_count >= min_valid_time).values,
        dims=("lat", "lon"),
        coords={"lat": obs_class["lat"], "lon": obs_class["lon"]},
        name="prob_metric_mask",
    )


# =========================================================
# Observation categorization
# =========================================================
def tercile_observations_from_training(
    y_obs: xr.DataArray,
    y_train: xr.DataArray,
    *,
    sample_dim: str = "time",
    use_safe_quantile: bool = True,
) -> xr.DataArray:
    """
    Categorize observations into terciles using thresholds from y_train.

    Output codes:
      0 = below
      1 = near
      2 = above
      NaN = invalid
    """
    _require_dims(y_obs, (sample_dim, "lat", "lon"), "y_obs")
    _require_dims(y_train, (sample_dim, "lat", "lon"), "y_train")

    if use_safe_quantile:
        q33 = _safe_nanquantile(y_train, 1.0 / 3.0, dim=sample_dim)
        q67 = _safe_nanquantile(y_train, 2.0 / 3.0, dim=sample_dim)
    else:
        q33 = y_train.quantile(1.0 / 3.0, dim=sample_dim)
        q67 = y_train.quantile(2.0 / 3.0, dim=sample_dim)

    valid = np.isfinite(y_obs) & np.isfinite(q33) & np.isfinite(q67)

    below = valid & (y_obs < q33)
    above = valid & (y_obs > q67)
    near = valid & (~below) & (~above)

    cls = xr.full_like(y_obs, fill_value=np.nan, dtype=float)
    cls = xr.where(below, 0.0, cls)
    cls = xr.where(near, 1.0, cls)
    cls = xr.where(above, 2.0, cls)

    return cls.rename("obs_tercile")


def one_hot_tercile(
    obs_class: xr.DataArray,
    *,
    sample_dim: str = "time",
    category_dim: str = "category",
    category_labels: Sequence[str] = DEFAULT_CATEGORY_LABELS,
) -> xr.DataArray:
    """
    Convert observation tercile classes (0/1/2) to one-hot probabilities.
    """
    _require_dims(obs_class, (sample_dim, "lat", "lon"), "obs_class")

    if len(category_labels) != 3:
        raise ValueError(f"category_labels must have length 3, got {category_labels}")

    arrs = []
    for code, lab in zip(DEFAULT_CATEGORY_CODES, category_labels):
        arr = (obs_class == code).astype(float).assign_coords({category_dim: lab})
        arrs.append(arr)

    return xr.concat(arrs, dim=category_dim).transpose(category_dim, sample_dim, "lat", "lon")


# =========================================================
# Ranked probability skill score
# =========================================================
def compute_rpss_map(
    prob_fcst: xr.DataArray,
    obs_class: xr.DataArray,
    *,
    category_dim: str = "category",
    sample_dim: str = "time",
    category_labels: Sequence[str] = DEFAULT_CATEGORY_LABELS,
    clip_probabilities: bool = False,
    normalize_probabilities: bool = False,
) -> xr.DataArray:
    """
    Compute gridpoint RPSS for 3-category forecasts.

    RPSS = 1 - mean(RPS_forecast) / mean(RPS_reference)

    The reference is climatology with equal 1/3 probability in each category.
    """
    _require_dims(obs_class, (sample_dim, "lat", "lon"), "obs_class")
    _validate_probabilities(prob_fcst, category_dim=category_dim, sample_dim=sample_dim)

    if len(category_labels) != 3:
        raise ValueError(f"category_labels must have length 3, got {category_labels}")

    fcst = prob_fcst.copy()

    if clip_probabilities:
        fcst = fcst.clip(min=0.0, max=1.0)

    if normalize_probabilities:
        denom = fcst.sum(dim=category_dim)
        denom = xr.where(np.abs(denom) < 1e-12, np.nan, denom)
        fcst = fcst / denom

    obs_oh = one_hot_tercile(
        obs_class,
        sample_dim=sample_dim,
        category_dim=category_dim,
        category_labels=category_labels,
    )

    fcst_cum = fcst.cumsum(dim=category_dim)
    obs_cum = obs_oh.cumsum(dim=category_dim)
    rps = ((fcst_cum - obs_cum) ** 2).sum(dim=category_dim)

    ref = xr.full_like(fcst, 1.0 / 3.0)
    ref_cum = ref.cumsum(dim=category_dim)
    rps_ref = ((ref_cum - obs_cum) ** 2).sum(dim=category_dim)

    rps_mean = rps.mean(dim=sample_dim, skipna=True)
    rps_ref_mean = rps_ref.mean(dim=sample_dim, skipna=True)

    rpss = 1.0 - (rps_mean / rps_ref_mean)
    rpss = xr.where(np.isfinite(rps_ref_mean) & (rps_ref_mean > 0), rpss, np.nan)

    return rpss.rename("rpss")


# =========================================================
# Deterministic metric
# =========================================================
def compute_correlation_map(
    y_pred: xr.DataArray,
    y_obs: xr.DataArray,
    *,
    sample_dim: str = "time",
    min_valid_time: int = 3,
) -> xr.DataArray:
    """
    Correlation map with paired-sample screening.
    """
    _require_dims(y_pred, (sample_dim, "lat", "lon"), "y_pred")
    _require_dims(y_obs, (sample_dim, "lat", "lon"), "y_obs")

    mask = build_metric_mask(
        y_obs,
        y_pred,
        sample_dim=sample_dim,
        min_valid_time=min_valid_time,
    )
    y_pred_m = y_pred.where(mask)
    y_obs_m = y_obs.where(mask)

    out = xr.corr(y_pred_m, y_obs_m, dim=sample_dim)
    return out.rename("correlation")


# =========================================================
# Probabilistic scores
# =========================================================
def compute_grocs_map(
    prob_fcst: xr.DataArray,
    obs_class: xr.DataArray,
    *,
    sample_dim: str = "time",
    category_dim: str = "category",
    min_valid_time: int = 3,
) -> xr.DataArray:
    """
    Practical simplified multi-category GROC skill:
    mean one-vs-rest AUC across the 3 tercile categories, rescaled to [-1, 1].
    """
    _require_dims(obs_class, (sample_dim, "lat", "lon"), "obs_class")
    _validate_probabilities(prob_fcst, category_dim=category_dim, sample_dim=sample_dim)

    label_below = _resolve_category_label(prob_fcst, "below", category_dim)
    label_near = _resolve_category_label(prob_fcst, "near", category_dim)
    label_above = _resolve_category_label(prob_fcst, "above", category_dim)

    auc_below = compute_binary_auc_map(
        prob_fcst.sel({category_dim: label_below}),
        (obs_class == 0),
        sample_dim=sample_dim,
        min_valid_time=min_valid_time,
    )
    auc_near = compute_binary_auc_map(
        prob_fcst.sel({category_dim: label_near}),
        (obs_class == 1),
        sample_dim=sample_dim,
        min_valid_time=min_valid_time,
    )
    auc_above = compute_binary_auc_map(
        prob_fcst.sel({category_dim: label_above}),
        (obs_class == 2),
        sample_dim=sample_dim,
        min_valid_time=min_valid_time,
    )

    auc_mean = xr.concat([auc_below, auc_near, auc_above], dim="tmp").mean("tmp", skipna=True)
    grocs = 2.0 * auc_mean - 1.0
    return grocs.rename("grocs")


def compute_binary_auc_map(
    prob: xr.DataArray,
    event: xr.DataArray,
    *,
    sample_dim: str = "time",
    min_valid_time: int = 3,
) -> xr.DataArray:
    """
    Rank-based AUC over time, gridpoint by gridpoint.

    Uses average ranks for tied probabilities.
    """
    if min_valid_time < 1:
        raise ValueError("min_valid_time must be >= 1")

    _require_dims(prob, (sample_dim, "lat", "lon"), "prob")
    _require_dims(event, (sample_dim, "lat", "lon"), "event")

    p = prob.transpose(sample_dim, "lat", "lon").values
    e = event.transpose(sample_dim, "lat", "lon").values

    _, ny, nx = p.shape
    out = np.full((ny, nx), np.nan, dtype=float)

    for j in range(ny):
        for i in range(nx):
            pj = np.asarray(p[:, j, i], dtype=float)
            ej = np.asarray(e[:, j, i])

            valid = np.isfinite(pj) & np.isfinite(ej)
            pj = pj[valid]
            ej = ej[valid].astype(bool)

            if pj.size < min_valid_time:
                continue

            n1 = int(np.sum(ej))
            n0 = int(np.sum(~ej))
            if n1 < 2 or n0 < 2:
                continue

            order = np.argsort(pj, kind="mergesort")
            sorted_p = pj[order]

            ranks_sorted = np.empty_like(sorted_p, dtype=float)
            start = 0
            n = len(sorted_p)
            while start < n:
                end = start + 1
                while end < n and sorted_p[end] == sorted_p[start]:
                    end += 1
                avg_rank = 0.5 * (start + 1 + end)
                ranks_sorted[start:end] = avg_rank
                start = end

            ranks = np.empty_like(ranks_sorted, dtype=float)
            ranks[order] = ranks_sorted

            rank_sum_pos = np.sum(ranks[ej])
            auc = (rank_sum_pos - n1 * (n1 + 1) / 2.0) / (n1 * n0)
            out[j, i] = auc

    return xr.DataArray(
        out,
        dims=("lat", "lon"),
        coords={"lat": prob["lat"], "lon": prob["lon"]},
        name="auc",
    )


# =========================================================
# Domain summaries
# =========================================================
def compute_domain_roc_curves(
    prob_fcst: xr.DataArray,
    obs_class: xr.DataArray,
    *,
    category_dim: str = "category",
    n_thresh: int = 11,
) -> dict:
    """
    Domain-flattened ROC curves for each tercile class.
    """
    if n_thresh < 2:
        raise ValueError("n_thresh must be at least 2")

    _require_dims(obs_class, ("time", "lat", "lon"), "obs_class")
    _validate_probabilities(prob_fcst, category_dim=category_dim, sample_dim="time")

    thresholds = np.linspace(0.0, 1.0, n_thresh)
    out = {}

    for label, cls in zip(DEFAULT_CATEGORY_LABELS, DEFAULT_CATEGORY_CODES):
        resolved_label = _resolve_category_label(prob_fcst, label, category_dim)
        p = np.asarray(prob_fcst.sel({category_dim: resolved_label}).values, dtype=float).ravel()
        o = np.asarray(obs_class.values == cls, dtype=bool).ravel()

        valid = np.isfinite(p)
        p = p[valid]
        o = o[valid]

        fpr = []
        tpr = []
        for thr in thresholds:
            yes = p >= thr
            tp = np.sum(yes & o)
            fp = np.sum(yes & (~o))
            fn = np.sum((~yes) & o)
            tn = np.sum((~yes) & (~o))

            tpr.append(tp / (tp + fn) if (tp + fn) > 0 else np.nan)
            fpr.append(fp / (fp + tn) if (fp + tn) > 0 else np.nan)

        out[label] = {
            "thresholds": thresholds,
            "fpr": np.asarray(fpr, dtype=float),
            "tpr": np.asarray(tpr, dtype=float),
        }

    return out


def compute_domain_reliability(
    prob_fcst: xr.DataArray,
    obs_class: xr.DataArray,
    *,
    category_dim: str = "category",
    n_bins: int = 10,
) -> dict:
    """
    Domain-flattened reliability data for each tercile class.
    """
    if n_bins < 1:
        raise ValueError("n_bins must be at least 1")

    _require_dims(obs_class, ("time", "lat", "lon"), "obs_class")
    _validate_probabilities(prob_fcst, category_dim=category_dim, sample_dim="time")

    edges = np.linspace(0.0, 1.0, n_bins + 1)
    centers = 0.5 * (edges[:-1] + edges[1:])
    out = {}

    for label, cls in zip(DEFAULT_CATEGORY_LABELS, DEFAULT_CATEGORY_CODES):
        resolved_label = _resolve_category_label(prob_fcst, label, category_dim)
        p = np.asarray(prob_fcst.sel({category_dim: resolved_label}).values, dtype=float).ravel()
        o = np.asarray(obs_class.values == cls, dtype=float).ravel()

        valid = np.isfinite(p) & np.isfinite(o)
        p = p[valid]
        o = o[valid]

        forecast_bin = np.full(n_bins, np.nan, dtype=float)
        observed_bin = np.full(n_bins, np.nan, dtype=float)
        counts = np.zeros(n_bins, dtype=int)

        for b in range(n_bins):
            if b < n_bins - 1:
                mask = (p >= edges[b]) & (p < edges[b + 1])
            else:
                mask = (p >= edges[b]) & (p <= edges[b + 1])

            counts[b] = int(np.sum(mask))
            if counts[b] > 0:
                forecast_bin[b] = float(np.mean(p[mask]))
                observed_bin[b] = float(np.mean(o[mask]))

        out[label] = {
            "bin_centers": centers,
            "forecast_bin_mean": forecast_bin,
            "observed_frequency": observed_bin,
            "counts": counts,
        }

    return out

def compute_hss_map_from_classes(
    pred_class: xr.DataArray,
    obs_class: xr.DataArray,
    *,
    sample_dim: str = "time",
    min_valid_time: int = 3,
    category_codes: Sequence[int] = DEFAULT_CATEGORY_CODES,
    baseline: Sequence[float] = (1 / 3, 1 / 3, 1 / 3),
) -> xr.DataArray:
    """
    Compute Heidke Skill Score from categorical forecasts.

    Matches the R implementation using:
      HSS = (PC - chance) / (1 - chance)

    where:
      PC = proportion correct
      chance = sum(p_base * p_forecast)
    """

    _require_dims(pred_class, (sample_dim, "lat", "lon"), "pred_class")
    _require_dims(obs_class, (sample_dim, "lat", "lon"), "obs_class")

    pred, obs = xr.align(pred_class, obs_class, join="inner")

    p = pred.transpose(sample_dim, "lat", "lon").values
    o = obs.transpose(sample_dim, "lat", "lon").values

    _, ny, nx = p.shape

    out = np.full((ny, nx), np.nan, dtype=float)

    codes = list(category_codes)
    base = np.asarray(baseline, dtype=float)

    for j in range(ny):
        for i in range(nx):

            pj = np.asarray(p[:, j, i], dtype=float)
            oj = np.asarray(o[:, j, i], dtype=float)

            valid = np.isfinite(pj) & np.isfinite(oj)

            pj = pj[valid]
            oj = oj[valid]

            if pj.size < min_valid_time:
                continue

            pj = pj.astype(int)
            oj = oj.astype(int)

            pc = np.mean(pj == oj)

            p_forecast = np.array(
                [np.mean(pj == c) for c in codes],
                dtype=float,
            )

            chance = np.sum(base * p_forecast)

            denom = 1.0 - chance

            if np.abs(denom) < 1e-12:
                continue

            out[j, i] = (pc - chance) / denom

    return xr.DataArray(
        out,
        dims=("lat", "lon"),
        coords={
            "lat": obs_class["lat"],
            "lon": obs_class["lon"],
        },
        name="hss",
    )

def compute_hss_map(
    prob_fcst: xr.DataArray,
    obs_class: xr.DataArray,
    *,
    category_dim: str = "category",
    sample_dim: str = "time",
    min_valid_time: int = 3,
) -> xr.DataArray:
    """
    Compute gridpoint Heidke Skill Score from probabilistic tercile forecasts.

    Uses dominant tercile category from maximum probability.
    Handles all-NaN probability slices safely.
    """

    _require_dims(obs_class, (sample_dim, "lat", "lon"), "obs_class")

    _validate_probabilities(
        prob_fcst,
        category_dim=category_dim,
        sample_dim=sample_dim,
    )

    # ---------------------------------------------------------
    # Fill all-NaN category slices temporarily to avoid argmax crash
    # ---------------------------------------------------------
    all_nan_mask = prob_fcst.isnull().all(dim=category_dim)

    prob_safe = prob_fcst.fillna(-9999.0)

    pred_class = prob_safe.argmax(dim=category_dim).astype(float)

    # restore invalid cells to NaN
    pred_class = pred_class.where(~all_nan_mask)

    pred_class = pred_class.rename("pred_tercile")

    return compute_hss_map_from_classes(
        pred_class,
        obs_class,
        sample_dim=sample_dim,
        min_valid_time=min_valid_time,
        category_codes=DEFAULT_CATEGORY_CODES,
    )
