from __future__ import annotations
from dataclasses import dataclass
import numpy as np
import xarray as xr

@dataclass
class SpeiFit:
    shape: np.ndarray
    loc: np.ndarray
    scale: np.ndarray
    timescales: np.ndarray
    months: np.ndarray


def rolling_sum(arr: np.ndarray, window: int) -> np.ndarray:
    x = np.asarray(arr, dtype=float)
    out = np.full_like(x, np.nan, dtype=np.float32)
    if int(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)
    pad0 = np.zeros_like(cs[:1])
    sums = cs[window-1:] - np.concatenate([pad0, cs[:-window]], axis=0)
    counts = cf[window-1:] - np.concatenate([pad0.astype(np.int16), cf[:-window]], axis=0)
    out[window-1:] = np.where(counts == window, sums, np.nan).astype(np.float32)
    return out


def fit_loglogistic(dp: np.ndarray, months: np.ndarray, ref_mask: np.ndarray, timescales=(1,3,6), min_samples=20) -> SpeiFit:
    from scipy.stats import fisk
    ts = np.asarray(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, w in enumerate(ts):
        acc = rolling_sum(dp, int(w))
        for m in range(1, 13):
            vals = acc[ref_mask & (months == m)]
            if vals.size == 0:
                continue
            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)
                        if np.isfinite(c) and np.isfinite(l) and np.isfinite(s) and c > 0 and s > 0:
                            shape[it,m-1,j,i] = c; loc[it,m-1,j,i] = l; scale[it,m-1,j,i] = s
                    except Exception:
                        pass
    return SpeiFit(shape, loc, scale, ts, np.arange(1,13,dtype=int))


def transform_spei(dp: np.ndarray, months: np.ndarray, fit: SpeiFit, cdf_eps=1e-6) -> dict[int, np.ndarray]:
    from scipy.stats import fisk, norm
    out = {}
    for it, w in enumerate(fit.timescales):
        acc = rolling_sum(dp, int(w))
        spei = np.full_like(acc, np.nan, dtype=np.float32)
        for m in range(1,13):
            kk = months == m
            if not np.any(kk):
                continue
            vals = acc[kk]
            c = fit.shape[it,m-1]; l = fit.loc[it,m-1]; s = fit.scale[it,m-1]
            valid_param = np.isfinite(c) & np.isfinite(l) & np.isfinite(s) & (c>0) & (s>0)
            ov = np.full_like(vals, np.nan, dtype=np.float32)
            for t in range(vals.shape[0]):
                x = vals[t]
                ok = valid_param & np.isfinite(x)
                p = np.full_like(x, np.nan, dtype=float)
                if np.any(ok):
                    p[ok] = fisk.cdf(x[ok], c[ok], loc=l[ok], scale=s[ok])
                    p = np.clip(p, cdf_eps, 1.0 - cdf_eps)
                    ov[t] = norm.ppf(p).astype(np.float32)
            spei[kk] = ov
        out[int(w)] = spei
    return out


def fit_to_dataset(fit: SpeiFit, lat, lon) -> xr.Dataset:
    return xr.Dataset(
        coords={'timescale': fit.timescales, 'calendar_month': fit.months, 'lat': lat, 'lon': lon},
        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': 'gridpoint_calendar_month_timescale'}
    )


def fit_from_dataset(ds: xr.Dataset) -> SpeiFit:
    return SpeiFit(ds['shape'].values, ds['loc'].values, ds['scale'].values, ds['timescale'].values.astype(int), ds['calendar_month'].values.astype(int))