"""Foster & Rahmstorf's adjustment (Xogen.r, Zenodo 15608929) and Beaulieu et al.'s PLF-changepoint
test, reimplemented in numpy so the whole pipeline can be run thousands of times."""
import numpy as np
from statsmodels.nonparametric.smoothers_lowess import lowess

K_VOLC, K_NINO, K_SSN = 25, 13, 13  # lags 0..24, 0..12, 0..12 months, as in Xogen.r


def lowess_r(t, x, frac):
    """R's lowess(t, x, f=frac): 3 robustness iterations, delta = 1% of the range of t."""
    return lowess(x, t, frac=frac, it=3, delta=0.01 * (t.max() - t.min()), return_sorted=False)


def lagged(x, k):
    """Xogen.r's simplag(method="lag"): shift by k months, padding with the first value."""
    if k == 0:
        return x.copy()
    return np.concatenate([np.full(k, x[0]), x[:-k]])


class Adjuster:
    """The exogenous fit for one series' time axis and predictors. Everything that doesn't depend on the
    temperatures is computed once, so each pass of the lag search costs a few dot products."""

    def __init__(self, t, volc, nino, ssn, fit=(1950, 2025), nback=4):
        self.t, self.nback = t, nback
        self.frac = 20 / (t.max() - t.min())
        self.win = (t >= fit[0]) & (t < fit[1])
        self.L = [np.stack([lagged(s, k) for k in range(n)]) for s, n in ((volc, K_VOLC), (nino, K_NINO), (ssn, K_SSN))]
        W = [m[:, self.win] for m in self.L]
        self.Wc = [m - m.mean(axis=1, keepdims=True) for m in W]  # centred within the fit window
        V, N, S = self.Wc
        G = np.empty((K_VOLC, K_NINO, K_SSN, 3, 3))
        G[..., 0, 0] = (V * V).sum(1)[:, None, None]
        G[..., 1, 1] = (N * N).sum(1)[None, :, None]
        G[..., 2, 2] = (S * S).sum(1)[None, None, :]
        G[..., 0, 1] = G[..., 1, 0] = (V @ N.T)[:, :, None]
        G[..., 0, 2] = G[..., 2, 0] = (V @ S.T)[:, None, :]
        G[..., 1, 2] = G[..., 2, 1] = (N @ S.T)[None, :, :]
        self.Ginv = np.linalg.inv(G)

    def fit_pass(self, x, y):
        """fit.xogen: choose the lags whose regression of y on the lagged predictors has the lowest AIC
        (with the same number of parameters for every choice, the lowest residual sum of squares), then
        return the adjusted series x minus the full fit, rounded to 4 decimals as Xogen.r rounds it."""
        yw = y[self.win] - y[self.win].mean()
        c = np.stack(np.broadcast_arrays((self.Wc[0] @ yw)[:, None, None], (self.Wc[1] @ yw)[None, :, None],
                                         (self.Wc[2] @ yw)[None, None, :]), axis=-1)
        explained = np.einsum("...i,...ij,...j->...", c, self.Ginv, c)
        k = np.unravel_index(np.argmax(explained), explained.shape)  # first maximum in R's loop order
        X = np.column_stack([np.ones(self.win.sum())] + [self.L[j][k[j]][self.win] for j in range(3)])
        coef = np.linalg.lstsq(X, y[self.win], rcond=None)[0]
        fitted = coef[0] + sum(coef[j + 1] * self.L[j][k[j]] for j in range(3))
        return np.round(x - fitted, 4), k, coef

    def adjust(self, x):
        """xogen(): lowess first pass, then nback rounds of backfitting."""
        z, k, coef = self.fit_pass(x, x - lowess_r(self.t, x, self.frac))
        for _ in range(self.nback):
            z, k, coef = self.fit_pass(x, x - lowess_r(self.t, z, self.frac))
        return z, k, coef


def annual(t, z, start, end):
    """Calendar-year means of monthly values, for years start..end."""
    yr = np.floor(t).astype(int)
    years = np.arange(start, end + 1)
    return years, np.array([z[yr == y].mean() for y in years])