"""Leakage-safe verification primitives for continuous and drought forecasts."""
from __future__ import annotations

from collections.abc import Iterable

import numpy as np
import xarray as xr


def _dimensions(dim: str | Iterable[str]) -> list[str]:
    return [dim] if isinstance(dim, str) else list(dim)


def _aligned(forecast: xr.DataArray, observation: xr.DataArray) -> tuple[xr.DataArray, xr.DataArray]:
    left, right = xr.align(forecast, observation, join="exact")
    return left.astype(float), right.astype(float)


def continuous_metrics(forecast: xr.DataArray, observation: xr.DataArray,
                       dim: str | Iterable[str], precipitation: bool = False) -> xr.Dataset:
    """Compute confirmed continuous metrics over explicitly supplied sample dimensions.

    Temperature returns ME, MAE, RMSE and Pearson correlation. Precipitation
    returns ME, RMAE and correlation, where RMAE=MAE/mean(observation)*100.
    Only paired finite samples contribute, and their count is returned.
    """
    dimensions = _dimensions(dim)
    missing = [name for name in dimensions if name not in forecast.dims or name not in observation.dims]
    if missing:
        raise ValueError(f"Verification dimensions missing from inputs: {missing}")
    model, obs = _aligned(forecast, observation)
    paired = np.isfinite(model) & np.isfinite(obs)
    model = model.where(paired)
    obs = obs.where(paired)
    error = model - obs
    count = paired.sum(dimensions).rename("sample_count")
    me = error.mean(dimensions, skipna=True).rename("me")
    model_anomaly = model - model.mean(dimensions, skipna=True)
    obs_anomaly = obs - obs.mean(dimensions, skipna=True)
    covariance_sum = (model_anomaly * obs_anomaly).sum(dimensions, skipna=True)
    spread = np.sqrt(
        (model_anomaly**2).sum(dimensions, skipna=True)
        * (obs_anomaly**2).sum(dimensions, skipna=True)
    )
    correlation = (covariance_sum / spread).where((count >= 2) & (spread > 0)).rename("correlation")
    output = {"me": me, "correlation": correlation, "sample_count": count}
    if precipitation:
        mae = abs(error).mean(dimensions, skipna=True)
        mean_obs = obs.mean(dimensions, skipna=True)
        rmae = (100.0 * mae / mean_obs).where(mean_obs > 0).rename("rmae")
        rmae.attrs = {"units": "%", "definition": "MAE / mean(paired observations) * 100"}
        output["rmae"] = rmae
    else:
        output["mae"] = abs(error).mean(dimensions, skipna=True).rename("mae")
        output["rmse"] = np.sqrt((error**2).mean(dimensions, skipna=True)).rename("rmse")
    result = xr.Dataset(output)
    result.attrs = {"pairing": "only paired finite forecast-observation samples",
                    "selection_note": "compute separately by lead, calendar month/season and region upstream"}
    return result


def spei_drought_event(spei: xr.DataArray, level: int) -> xr.DataArray:
    """Create D1..D4 binary events using the confirmed SPEI thresholds."""
    if level not in (1, 2, 3, 4):
        raise ValueError("Drought level must be one of 1, 2, 3, 4")
    threshold = {1: -0.5, 2: -1.0, 3: -1.5, 4: -2.0}[level]
    event = (spei < threshold) if level == 4 else (spei <= threshold)
    result = event.astype(float).where(np.isfinite(spei)).rename(f"d{level}_event")
    result.attrs = {"units": "1", "threshold": f"SPEI {'<' if level == 4 else '<='} {threshold}"}
    return result


def contingency_metrics(forecast_event: xr.DataArray, observed_event: xr.DataArray,
                        dim: str | Iterable[str]) -> xr.Dataset:
    """Compute POD, POFD, FAR, FBI, CSI, EDI and SEDI without edge correction."""
    dimensions = _dimensions(dim)
    forecast, observed = _aligned(forecast_event, observed_event)
    paired = np.isfinite(forecast) & np.isfinite(observed)
    invalid = paired & ~forecast.isin([0.0, 1.0])
    invalid = invalid | (paired & ~observed.isin([0.0, 1.0]))
    if bool(invalid.any().item()):
        raise ValueError("Event inputs must contain only 0, 1 or missing")
    yes_f = forecast == 1
    yes_o = observed == 1
    hits = (paired & yes_f & yes_o).sum(dimensions).rename("hits")
    misses = (paired & ~yes_f & yes_o).sum(dimensions).rename("misses")
    false_alarms = (paired & yes_f & ~yes_o).sum(dimensions).rename("false_alarms")
    correct_negatives = (paired & ~yes_f & ~yes_o).sum(dimensions).rename("correct_negatives")

    def ratio(numerator: xr.DataArray, denominator: xr.DataArray) -> xr.DataArray:
        return (numerator / denominator).where(denominator > 0)

    pod = ratio(hits, hits + misses).rename("pod")
    pofd = ratio(false_alarms, false_alarms + correct_negatives).rename("pofd")
    far = ratio(false_alarms, hits + false_alarms).rename("far")
    fbi = ratio(hits + false_alarms, hits + misses).rename("fbi")
    csi = ratio(hits, hits + misses + false_alarms).rename("csi")
    interior = (pod > 0) & (pod < 1) & (pofd > 0) & (pofd < 1)
    safe_pod = pod.where(interior, 0.5)
    safe_pofd = pofd.where(interior, 0.5)
    edi = ((np.log(safe_pofd) - np.log(safe_pod)) /
           (np.log(safe_pofd) + np.log(safe_pod))).where(interior).rename("edi")
    sedi = ((np.log(safe_pofd) - np.log(safe_pod) - np.log1p(-safe_pofd) + np.log1p(-safe_pod)) /
            (np.log(safe_pofd) + np.log(safe_pod) + np.log1p(-safe_pofd) + np.log1p(-safe_pod))).where(interior).rename("sedi")
    result = xr.Dataset({array.name: array for array in (
        hits, misses, false_alarms, correct_negatives, pod, pofd, far, fbi, csi, edi, sedi
    )})
    result.attrs = {"missing_policy": "paired finite samples only",
                    "edi_sedi_edge_policy": "undefined rates 0 or 1 return missing; no correction applied"}
    return result