"""Beaulieu et al.'s test for a change in warming rate at an unknown time (Communications Earth &
Environment 2024, Methods "Differences of slopes test"), as Foster & Rahmstorf 2026 apply it:
annual values, a continuous piecewise-linear trend with one changepoint k, Gaussian AR(1) errors
fitted by exact maximum likelihood for each k, T_k = slope change / its standard error, and
T_max = max |T_k| over 0.1 N <= k <= 0.9 N. Critical values come from Monte Carlo AR(1) series.

Every sum the AR(1) likelihood needs is a quadratic polynomial in phi, so a batch of series is fitted
at many phi values with 3 x 3 solves only."""
import numpy as np

PHI_COARSE = np.linspace(-0.95, 0.95, 21)
REFINE_ROUNDS = 3


def ks(n):
    return np.arange(int(np.ceil(0.1 * n)), int(np.floor(0.9 * n)) + 1)


def design(n):
    t = np.arange(1, n + 1, dtype=float)
    K = ks(n)
    X = np.empty((len(K), n, 3))
    X[:, :, 0] = 1.0
    X[:, :, 1] = t
    X[:, :, 2] = np.maximum(t[None, :] - K[:, None], 0.0)
    return K, X


def _parts(a, b):
    """For column blocks a (..., n, p) and b (..., n, q): the pieces of sum_t a~_t b~_t' as a polynomial
    in phi, where a~_1 = sqrt(1 - phi^2) a_1 and a~_t = a_t - phi a_{t-1}."""
    first = a[..., 0, :, None] * b[..., 0, None, :]
    s0 = np.einsum("...ti,...tj->...ij", a[..., 1:, :], b[..., 1:, :])
    s1 = np.einsum("...ti,...tj->...ij", a[..., 1:, :], b[..., :-1, :]) + \
        np.einsum("...ti,...tj->...ij", a[..., :-1, :], b[..., 1:, :])
    s2 = np.einsum("...ti,...tj->...ij", a[..., :-1, :], b[..., :-1, :])
    return first, s0, s1, s2


def inv3(A):
    """Inverses of many symmetric 3 x 3 matrices, by cofactors."""
    a, b, c = A[..., 0, 0], A[..., 0, 1], A[..., 0, 2]
    d, e, f = A[..., 1, 1], A[..., 1, 2], A[..., 2, 2]
    c00, c01, c02 = d * f - e * e, c * e - b * f, b * e - c * d
    c11, c12, c22 = a * f - c * c, b * c - a * e, a * d - b * b
    det = a * c00 + b * c01 + c * c02
    out = np.stack([c00, c01, c02, c01, c11, c12, c02, c12, c22], axis=-1) / det[..., None]
    return out.reshape(A.shape)


def _at(parts, phi):
    first, s0, s1, s2 = parts
    p = phi.reshape(phi.shape + (1, 1))
    return (1 - p * p) * first + s0 - p * s1 + p * p * s2


class TmaxTest:
    def __init__(self, n):
        self.n = n
        self.K, X = design(n)
        self.XX = _parts(X, X)  # (K, 3, 3) each
        self.X = X

    def _fit(self, XX, Xy, yy, phi):
        """GLS at phi (shape (B, K, m)): profile log-likelihood, slope change, and its standard error."""
        A = _at(tuple(np.broadcast_to(p[None, :, None], phi.shape + (3, 3)) for p in XX), phi)
        c = _at(tuple(np.broadcast_to(p[:, :, None], phi.shape + (3, 1)) for p in Xy), phi)[..., 0]
        s = _at(tuple(np.broadcast_to(p.reshape(-1, 1, 1, 1, 1), phi.shape + (1, 1)) for p in yy), phi)[..., 0, 0]
        Ainv = inv3(A)
        beta = np.einsum("...ij,...j->...i", Ainv, c)
        rss = np.maximum(s - np.einsum("...i,...i->...", beta, c), 1e-300)
        ll = -0.5 * self.n * np.log(rss / self.n) + 0.5 * np.log(1 - phi * phi)
        se = np.sqrt(rss / self.n * Ainv[..., 2, 2])
        return ll, beta[..., 2], se

    def stats(self, y):
        """y: (B, n). Returns T_max (B,), the k it occurs at (B,), and T_k (B, K)."""
        y = np.atleast_2d(np.asarray(y, float))
        B = y.shape[0]
        Xy = _parts(self.X[None, :, :, :], np.broadcast_to(y[:, None, :, None], (B, len(self.K), self.n, 1)))
        yy = tuple(p[..., 0, 0] for p in _parts(y[:, :, None], y[:, :, None]))
        # a coarse grid, three rounds of finer grids around the best value, then a parabola through the best three
        K = len(self.K)
        phi = np.broadcast_to(PHI_COARSE, (B, K, len(PHI_COARSE)))
        ll, _, _ = self._fit(self.XX, Xy, yy, phi)
        best = PHI_COARSE[np.argmax(ll, axis=-1)]
        step = PHI_COARSE[1] - PHI_COARSE[0]
        offs = np.arange(-3, 4)
        for _ in range(REFINE_ROUNDS):
            step /= 3
            phi = np.clip(best[..., None] + step * offs, -0.995, 0.995)
            ll, _, _ = self._fit(self.XX, Xy, yy, phi)
            j = np.argmax(ll, axis=-1)
            best = np.take_along_axis(phi, j[..., None], -1)[..., 0]
        j = np.clip(j, 1, len(offs) - 2)
        l0, l1, l2 = (np.take_along_axis(ll, (j + d)[..., None], -1)[..., 0] for d in (-1, 0, 1))
        den = l0 - 2 * l1 + l2
        shift = np.where(den < 0, 0.5 * step * (l0 - l2) / np.where(den < 0, den, -1), 0.0)
        phi_hat = np.clip(np.take_along_axis(phi, j[..., None], -1)[..., 0] + np.clip(shift, -step, step), -0.995, 0.995)
        _, d, se = self._fit(self.XX, Xy, yy, phi_hat[..., None])
        T = (d / se)[..., 0]
        i = np.argmax(np.abs(T), axis=1)
        return np.abs(T)[np.arange(B), i], self.K[i], T


def ar1_null_fit(y):
    """The null model's AR(1) coefficient: exact Gaussian ML for a linear trend plus AR(1) errors."""
    n = len(y)
    t = np.arange(1, n + 1, dtype=float)
    X = np.column_stack([np.ones(n), t])[None]
    XX = _parts(X, X)
    Xy = _parts(X, y[None, :, None])
    yy = tuple(p[..., 0, 0] for p in _parts(y[None, :, None], y[None, :, None]))
    grid = np.linspace(-0.995, 0.995, 3981)

    def ll(phi):
        A = _at(tuple(np.broadcast_to(p, phi.shape + p.shape[-2:]) for p in XX), phi)
        c = _at(tuple(np.broadcast_to(p, phi.shape + p.shape[-2:]) for p in Xy), phi)[..., 0]
        s = _at(tuple(np.broadcast_to(np.asarray(p).reshape(1, 1, 1), phi.shape + (1, 1)) for p in yy), phi)[..., 0, 0]
        b = np.linalg.solve(A, c[..., None])[..., 0]
        rss = s - np.einsum("...i,...i->...", b, c)
        return -0.5 * n * np.log(rss / n) + 0.5 * np.log(1 - phi * phi), b, rss
    l, b, rss = ll(grid)
    i = int(np.argmax(l))
    return float(grid[i]), b[i], float(np.sqrt(rss[i] / n))


def simulate_ar1(rng, phi, n, size):
    """Stationary Gaussian AR(1) series with unit innovation variance."""
    e = rng.standard_normal((size, n))
    x = np.empty((size, n))
    x[:, 0] = e[:, 0] / np.sqrt(1 - phi * phi)
    for i in range(1, n):
        x[:, i] = phi * x[:, i - 1] + e[:, i]
    return x