ARIMA / SARIMAX

"""ARIMA by conditional sum of squares - the recursion, transcribed.

The likelihood reduces to the sum of squared one-step errors, and those errors
are produced by a recursion in the parameters rather than observed. That is the
whole difficulty: epsilon_t depends on epsilon_{t-1}, so there is no closed
form and the fit has to be numerical.
"""


def difference(series, d):
    """Apply (1 - B)^d. The I in ARIMA, and it happens before anything else."""
    result = list(series)
    for _ in range(d):
        result = [result[t] - result[t - 1] for t in range(1, len(result))]
    return result


def residuals(w, phi, theta):
    """epsilon_t = w_t - sum phi_i w_{t-i} - sum theta_j epsilon_{t-j}

    'Conditional' means the pre-sample values are assumed zero. That is the
    approximation this stage makes; the exact likelihood estimates them via a
    Kalman filter instead, which matters on short series.
    """
    p = len(phi)
    q = len(theta)
    errors = [0.0] * len(w)

    for t in range(len(w)):
        prediction = 0.0
        for i in range(p):
            if t - i - 1 >= 0:
                prediction += phi[i] * w[t - i - 1]
        for j in range(q):
            if t - j - 1 >= 0:
                prediction += theta[j] * errors[t - j - 1]
        errors[t] = w[t] - prediction

    return errors


def sum_of_squares(w, phi, theta):
    return sum(error * error for error in residuals(w, phi, theta))


def fit(series, p, d, q, lr=1e-4, iterations=2000, epsilon=1e-5):
    """Fit by gradient descent with numerical gradients.

    A real implementation uses L-BFGS with analytic derivatives; finite
    differences are used here because they show that the objective is just a
    function of the coefficients, with no structure the optimizer exploits.
    """
    w = difference(series, d)
    phi = [0.0] * p
    theta = [0.0] * q

    for _ in range(iterations):
        base = sum_of_squares(w, phi, theta)

        for i in range(p):
            phi[i] += epsilon
            gradient = (sum_of_squares(w, phi, theta) - base) / epsilon
            phi[i] -= epsilon
            phi[i] -= lr * gradient

        for j in range(q):
            theta[j] += epsilon
            gradient = (sum_of_squares(w, phi, theta) - base) / epsilon
            theta[j] -= epsilon
            theta[j] -= lr * gradient

    return phi, theta


def forecast(series, phi, theta, d, horizon):
    """Iterate the recursion forward, then undo the differencing."""
    w = difference(series, d)
    errors = residuals(w, phi, theta)

    history = list(w)
    future_errors = list(errors)

    for _ in range(horizon):
        prediction = 0.0
        for i in range(len(phi)):
            prediction += phi[i] * history[len(history) - i - 1]
        for j in range(len(theta)):
            prediction += theta[j] * future_errors[len(future_errors) - j - 1]
        history.append(prediction)
        # Future errors are zero in expectation - which is exactly why the
        # forecast reverts to the mean once the MA memory runs out.
        future_errors.append(0.0)

    # Undo the differencing by cumulative summation from the last observed
    # levels. Getting these starting values wrong shifts every forecast by a
    # constant that nobody notices.
    forecasts = history[len(w):]
    for _ in range(d):
        anchor = series[-1]
        forecasts = [anchor := anchor + step for step in forecasts]

    return forecasts

"""ARIMA - typed, constrained, fitted by L-BFGS on the exact CSS objective."""

from dataclasses import dataclass

import numpy as np
from numpy.typing import NDArray
from scipy.optimize import minimize

Vector = NDArray[np.float64]


@dataclass(frozen=True)
class ArimaOrder:
    p: int
    d: int
    q: int

    def __post_init__(self) -> None:
        if min(self.p, self.d, self.q) < 0:
            raise ValueError(f"orders must be non-negative, got {self}")
        if self.d > 2:
            raise ValueError(
                f"d={self.d}: differencing more than twice almost always means a "
                "transform is needed rather than another difference"
            )


@dataclass(frozen=True)
class ArimaModel:
    order: ArimaOrder
    phi: Vector
    theta: Vector
    sigma2: float
    anchors: Vector          # last d levels, needed to undo the differencing
    aic: float

    @property
    def is_stationary(self) -> bool:
        """AR polynomial roots must lie OUTSIDE the unit circle.

        A fitted model that fails this is not merely inaccurate - its
        forecasts diverge with horizon, and the optimizer will happily return
        one because the objective does not know about the constraint.
        """
        return _roots_outside_unit_circle(self.phi)

    @property
    def is_invertible(self) -> bool:
        """Same condition on the MA polynomial. A non-invertible model has an
        equivalent invertible representation, so the coefficients are not
        identified even though the fit is."""
        return _roots_outside_unit_circle(-self.theta)


def _roots_outside_unit_circle(coefficients: Vector) -> bool:
    if coefficients.size == 0:
        return True
    # 1 - c1 B - c2 B^2 ... ; numpy wants highest power first.
    polynomial = np.concatenate([[1.0], -coefficients])[::-1]
    return bool(np.all(np.abs(np.roots(polynomial)) > 1.0))


def _difference(series: Vector, d: int) -> Vector:
    return series if d == 0 else _difference(np.diff(series), d - 1)


def _residuals(w: Vector, phi: Vector, theta: Vector) -> Vector:
    """One pass of the recursion. Sequential in t by construction: the MA term
    reads an error the same loop is producing, so this cannot be vectorized
    over time no matter how it is written."""
    p, q = phi.size, theta.size
    errors = np.zeros(w.size, dtype=np.float64)

    for t in range(w.size):
        prediction = 0.0
        if p:
            back = min(p, t)
            prediction += float(phi[:back] @ w[t - 1 : t - back - 1 : -1]) if back else 0.0
        if q:
            back = min(q, t)
            prediction += float(theta[:back] @ errors[t - 1 : t - back - 1 : -1]) if back else 0.0
        errors[t] = w[t] - prediction

    return errors


def fit(series: Vector, order: ArimaOrder) -> ArimaModel:
    """Fit by conditional maximum likelihood. Raises ValueError on bad input."""
    if series.ndim != 1:
        raise ValueError(f"series must be 1-D, got shape {series.shape}")
    if series.size <= order.d + order.p + order.q + 1:
        raise ValueError(f"series of length {series.size} is too short for {order}")
    if not np.isfinite(series).all():
        raise ValueError("series contains NaN or inf; the recursion has no notion of missing")

    w = _difference(series, order.d)

    def objective(packed: Vector) -> float:
        phi, theta = packed[: order.p], packed[order.p :]
        return float(np.sum(_residuals(w, phi, theta) ** 2))

    start = np.zeros(order.p + order.q, dtype=np.float64)
    result = minimize(objective, start, method="L-BFGS-B",
                      bounds=[(-0.99, 0.99)] * start.size)

    phi, theta = result.x[: order.p], result.x[order.p :]
    errors = _residuals(w, phi, theta)
    sigma2 = float(np.mean(errors**2))

    # AIC over the DIFFERENCED series length: comparing models with different
    # d on different sample sizes is a common and invalid comparison.
    n_params = order.p + order.q + 1
    aic = w.size * np.log(sigma2) + 2 * n_params

    return ArimaModel(
        order=order,
        phi=phi,
        theta=theta,
        sigma2=sigma2,
        anchors=series[-max(order.d, 1) :],
        aic=float(aic),
    )

"""ARIMA - fused AR term, reused buffers, order grid in one pass."""

import numpy as np
from numpy.typing import NDArray
from scipy.optimize import minimize

Vector = NDArray[np.float64]
Matrix = NDArray[np.float64]


class ConditionalArima:
    """The AR contribution is a convolution; only the MA part is sequential.

    The recursion cannot be vectorized over time - the MA term reads an error
    the loop is still producing. But the AR term reads only OBSERVED values,
    so the whole AR contribution for every t is one matrix-vector product
    against a lag matrix built once. That leaves a scalar accumulation for the
    MA part, which is genuinely irreducible.

    The other saving is larger and less obvious: L-BFGS calls the objective
    tens of times per fit and hundreds across an order grid, so the residual
    buffer and the lag matrix are built ONCE per series rather than per call.
    """

    def __init__(self, series: Vector, max_lag: int) -> None:
        # One contiguous float64 block; every objective evaluation walks it.
        self._series = np.ascontiguousarray(series, dtype=np.float64)
        self._max_lag = max_lag
        self._cache: dict[int, tuple[Vector, Matrix]] = {}
        self._errors: Vector | None = None

    def _prepared(self, d: int) -> tuple[Vector, Matrix]:
        """Differenced series and its lag matrix, memoized per d.

        An order grid evaluates many (p, q) at the same d, so differencing and
        building the lag matrix once per d rather than once per candidate is
        most of the grid's cost removed.
        """
        if d not in self._cache:
            w = self._series
            for _ in range(d):
                w = np.diff(w)
            w = np.ascontiguousarray(w)

            lags = np.zeros((w.size, self._max_lag), dtype=np.float64)
            for lag in range(1, self._max_lag + 1):
                lags[lag:, lag - 1] = w[:-lag]
            self._cache[d] = (w, lags)

        return self._cache[d]

    def _sum_of_squares(self, w: Vector, lags: Matrix, phi: Vector, theta: Vector) -> float:
        # The entire AR contribution, for every t at once: one GEMV.
        ar_term = lags[:, : phi.size] @ phi if phi.size else np.zeros(w.size)

        if self._errors is None or self._errors.size != w.size:
            self._errors = np.empty(w.size, dtype=np.float64)   # allocated once
        errors = self._errors

        q = theta.size
        if q == 0:
            np.subtract(w, ar_term, out=errors)
            return float(errors @ errors)

        # Only the MA recursion remains sequential, and it is O(q) per step
        # rather than O(p + q).
        for t in range(w.size):
            ma_term = 0.0
            for j in range(min(q, t)):
                ma_term += theta[j] * errors[t - j - 1]
            errors[t] = w[t] - ar_term[t] - ma_term

        return float(errors @ errors)

    def fit_order(self, p: int, d: int, q: int) -> tuple[Vector, float]:
        w, lags = self._prepared(d)

        def objective(packed: Vector) -> float:
            return self._sum_of_squares(w, lags, packed[:p], packed[p:])

        result = minimize(
            objective,
            np.zeros(p + q),
            method="L-BFGS-B",
            bounds=[(-0.99, 0.99)] * (p + q),
        )

        sigma2 = result.fun / w.size
        aic = w.size * np.log(sigma2) + 2 * (p + q + 1)
        return result.x, float(aic)

    def select(self, max_p: int, max_d: int, max_q: int) -> tuple[tuple[int, int, int], Vector]:
        """AIC grid search. The shared per-d preparation is what makes this
        affordable - without it, the grid re-differences the series and
        rebuilds the lag matrix for every candidate."""
        best_order = (0, 0, 0)
        best_params = np.zeros(0)
        best_aic = np.inf

        for d in range(max_d + 1):
            for p in range(max_p + 1):
                for q in range(max_q + 1):
                    if p == 0 and q == 0:
                        continue
                    params, aic = self.fit_order(p, d, q)
                    if aic < best_aic:
                        best_aic, best_order, best_params = aic, (p, d, q), params

        return best_order, best_params