from __future__ import annotations

from dataclasses import dataclass
from typing import Dict, Iterable, Tuple

import numpy as np
import xarray as xr


@dataclass
class SpeiFitResult:
    shape: np.ndarray  # timescale, calendar_month, lat, lon
    loc: np.ndarray
    scale: np.ndarray
    timescales: np.ndarray
    calendar_months: np.ndarray


def rolling_accumulation(arr: np.ndarray, window: int) -> np.ndarray:
    """Moving accumulation along time. Any NaN in a window gives NaN."""
    x = np.asarray(arr, dtype=float)
    if x.ndim != 3:
        raise ValueError('arr must have shape (time, lat, lon)')
    if window < 1:
        raise ValueError('window must be >= 1')
    out = np.full_like(x, np.nan, dtype=np.float32)
    if window == 1:
        return x.astype(np.float32)
    finite = np.isfinite(x).astype(np.int16)
    x0 = np.where(np.isfinite(x), x, 0.0)
    cs = np.cumsum(x0, axis=0)
    cf = np.cumsum(finite, axis=0)
    sums = cs[window-1:] - np.concatenate([np.zeros_like(cs[:1]), cs[:-window]], axis=0)
    counts = cf[window-1:] - np.concatenate([np.zeros_like(cf[:1]), cf[:-window]], axis=0)
    out[window-1:] = np.where(counts == window, sums, np.nan).astype(np.float32)
    return out


def _require_scipy():
    try:
        from scipy.stats import fisk, norm  # noqa: F401
    except Exception as exc:  # pragma: no cover
        raise ImportError('SPEI fitting requires scipy. Install scipy on the login node before running SPEI steps.') from exc


def fit_loglogistic_3p_by_calendar_month(
    dp: np.ndarray,
    months: np.ndarray,
    reference_mask: np.ndarray,
    timescales: Iterable[int],
    min_samples: int = 20,
) -> SpeiFitResult:
    """Fit 3-parameter log-logistic distribution to accumulated Dp.

    Fit dimensions: timescale × calendar_month × gridpoint.
    Moving accumulation is computed on the full continuous monthly series first.
    """
    _require_scipy()
    from scipy.stats import fisk

    dp = np.asarray(dp, dtype=float)
    months = np.asarray(months, dtype=int)
    reference_mask = np.asarray(reference_mask, dtype=bool)
    if dp.ndim != 3:
        raise ValueError('dp must have shape (time, lat, lon)')
    if dp.shape[0] != months.size or months.size != reference_mask.size:
        raise ValueError('time dimension mismatch')
    ts = np.asarray(list(timescales), dtype=int)
    nlat, nlon = dp.shape[1], dp.shape[2]
    shape = np.full((len(ts), 12, nlat, nlon), np.nan, dtype=np.float32)
    loc = np.full_like(shape, np.nan)
    scale = np.full_like(shape, np.nan)

    for it, win in enumerate(ts):
        acc = rolling_accumulation(dp, int(win))
        for cm in range(1, 13):
            tmask = reference_mask & (months == cm)
            if not np.any(tmask):
                continue
            vals = acc[tmask]
            for j in range(nlat):
                for i in range(nlon):
                    x = vals[:, j, i]
                    x = x[np.isfinite(x)]
                    if x.size < min_samples or np.nanstd(x) <= 1e-8:
                        continue
                    try:
                        c, l, s = fisk.fit(x)  # 3-parameter log-logistic.
                        if np.isfinite(c) and np.isfinite(l) and np.isfinite(s) and c > 0 and s > 0:
                            shape[it, cm-1, j, i] = c
                            loc[it, cm-1, j, i] = l
                            scale[it, cm-1, j, i] = s
                    except Exception:
                        continue
    return SpeiFitResult(shape=shape, loc=loc, scale=scale, timescales=ts, calendar_months=np.arange(1, 13, dtype=int))


def transform_to_spei(
    dp: np.ndarray,
    months: np.ndarray,
    fit: SpeiFitResult,
    cdf_epsilon: float = 1e-6,
) -> Dict[int, np.ndarray]:
    """Transform Dp to SPEI for each fitted timescale."""
    _require_scipy()
    from scipy.stats import fisk, norm

    dp = np.asarray(dp, dtype=float)
    months = np.asarray(months, dtype=int)
    result: Dict[int, np.ndarray] = {}
    for it, win in enumerate(fit.timescales):
        acc = rolling_accumulation(dp, int(win))
        spei = np.full_like(acc, np.nan, dtype=np.float32)
        for cm in range(1, 13):
            tmask = months == cm
            if not np.any(tmask):
                continue
            c = fit.shape[it, cm-1]
            l = fit.loc[it, cm-1]
            s = fit.scale[it, cm-1]
            valid_param = np.isfinite(c) & np.isfinite(l) & np.isfinite(s) & (c > 0) & (s > 0)
            vals = acc[tmask]
            out_vals = np.full_like(vals, np.nan, dtype=np.float32)
            for kk in range(vals.shape[0]):
                x = vals[kk]
                valid = valid_param & np.isfinite(x)
                if not np.any(valid):
                    continue
                p = np.full_like(x, np.nan, dtype=float)
                p[valid] = fisk.cdf(x[valid], c[valid], loc=l[valid], scale=s[valid])
                p = np.clip(p, cdf_epsilon, 1.0 - cdf_epsilon)
                out_vals[kk] = norm.ppf(p).astype(np.float32)
            spei[tmask] = out_vals
        result[int(win)] = spei
    return result


def fit_result_to_dataset(fit: SpeiFitResult, lat: np.ndarray, lon: np.ndarray) -> xr.Dataset:
    coords = {'timescale': fit.timescales, 'calendar_month': fit.calendar_months, 'lat': lat, 'lon': lon}
    return xr.Dataset(
        coords=coords,
        data_vars={
            'shape': (('timescale', 'calendar_month', 'lat', 'lon'), fit.shape),
            'loc': (('timescale', 'calendar_month', 'lat', 'lon'), fit.loc),
            'scale': (('timescale', 'calendar_month', 'lat', 'lon'), fit.scale),
        },
        attrs={'distribution': 'three-parameter log-logistic', 'fit_by_calendar_month': 'true'},
    )


def fit_result_from_dataset(ds: xr.Dataset) -> SpeiFitResult:
    return SpeiFitResult(
        shape=ds['shape'].values.astype(np.float32),
        loc=ds['loc'].values.astype(np.float32),
        scale=ds['scale'].values.astype(np.float32),
        timescales=ds['timescale'].values.astype(int),
        calendar_months=ds['calendar_month'].values.astype(int),
    )