N-BEATS / N-HiTS

"""N-BEATS, transcribed the way the paper reads.

One block does four things:

    theta       = MLP(input)              four ReLU layers, nothing else
    backcast    = V_b @ theta_b           what this block claims to explain
    forecast    = V_f @ theta_f           what it predicts
    next_input  = input - backcast        the residual the next block sees

Forecasts ADD and inputs have their explained part SUBTRACTED. That pair of
recursions is the whole architecture, and it is gradient boosting rearranged
into a neural network.

Plain loops, no libraries.
"""

import math


def relu(value):
    return value if value > 0.0 else 0.0


def matvec(weight, vector, bias):
    """weight is out_dim rows of in_dim."""
    return [
        bias[i] + sum(w * v for w, v in zip(row, vector))
        for i, row in enumerate(weight)
    ]


def mlp(vector, layers):
    """Four fully connected ReLU layers. The entire nonlinearity in N-BEATS."""
    hidden = vector
    for weight, bias in layers:
        hidden = [relu(value) for value in matvec(weight, hidden, bias)]
    return hidden


def generic_basis(length):
    """Identity: theta IS the output, one coefficient per position.

    Maximum flexibility, zero structure, and no interpretability at all.
    """
    return [[1.0 if i == j else 0.0 for j in range(length)] for i in range(length)]


def trend_basis(length, degree):
    """Polynomial in normalized time: V[t][p] = (t / length)^p.

    This is the point of basis expansion. The MLP predicts a handful of
    polynomial coefficients rather than every output value, so the output is
    forced to be a smooth low-order curve. That constrains the solution far
    more effectively than any weight penalty, and it is why this stack can be
    labelled "trend" and plotted.
    """
    return [
        [(index / length) ** power for power in range(degree + 1)]
        for index in range(length)
    ]


def seasonality_basis(length, harmonics):
    """Fourier: alternating cosine and sine at increasing frequencies.

    Same idea as the trend basis with a periodic family instead. The MLP
    predicts amplitudes; the basis guarantees the output is periodic.
    """
    rows = []
    for index in range(length):
        row = []
        for harmonic in range(1, harmonics + 1):
            angle = 2.0 * math.pi * harmonic * index / length
            row.append(math.cos(angle))
            row.append(math.sin(angle))
        rows.append(row)
    return rows


def project(basis, theta):
    """basis is length rows of coefficient_count."""
    return [
        sum(b * t for b, t in zip(row, theta))
        for row in basis
    ]


def block_forward(window, layers, theta_backcast_weight, theta_forecast_weight,
                  backcast_basis, forecast_basis):
    """One block: shared trunk, two heads, two basis projections."""
    hidden = mlp(window, layers)

    # The two heads predict COEFFICIENTS, not values. With a generic basis
    # that distinction is vacuous; with trend or seasonality it is everything.
    theta_backcast = matvec(theta_backcast_weight, hidden,
                            [0.0] * len(theta_backcast_weight))
    theta_forecast = matvec(theta_forecast_weight, hidden,
                            [0.0] * len(theta_forecast_weight))

    backcast = project(backcast_basis, theta_backcast)
    forecast = project(forecast_basis, theta_forecast)
    return backcast, forecast


def nbeats_forward(window, blocks):
    """Doubly residual stacking: forecasts add, inputs shrink."""
    residual = list(window)
    horizon = len(blocks[0]['forecast_basis'])
    total_forecast = [0.0] * horizon

    for block in blocks:
        backcast, forecast = block_forward(
            residual,
            block['layers'],
            block['theta_backcast_weight'],
            block['theta_forecast_weight'],
            block['backcast_basis'],
            block['forecast_basis'],
        )
        # Subtract what this block explained. A later block that receives a
        # near-zero residual has nothing left to learn - which is exactly
        # how depth stops paying, and it is invisible unless you look.
        residual = [r - b for r, b in zip(residual, backcast)]
        total_forecast = [t + f for t, f in zip(total_forecast, forecast)]

    return total_forecast, residual


def mase(forecast, actual, history, seasonal_period):
    """Mean absolute scaled error.

    The denominator is the in-sample error of a seasonal-naive forecast, so
    the loss is comparable across series of wildly different scale. With a
    plain MSE a global model trained across a panel optimizes only the
    largest series and ignores everything else.
    """
    denominator = 0.0
    counted = 0
    for index in range(seasonal_period, len(history)):
        denominator += abs(history[index] - history[index - seasonal_period])
        counted += 1
    if counted == 0 or denominator == 0.0:
        raise ValueError("history is too short to scale the loss")
    denominator /= counted

    numerator = sum(abs(f - a) for f, a in zip(forecast, actual)) / len(forecast)
    return numerator / denominator


def ensemble_forecast(window, model_list):
    """Not a refinement. The published results average dozens of models
    trained with different lookback windows and losses, and a single model is
    materially worse - so the ensemble is part of the method."""
    horizon = None
    accumulated = None
    for blocks in model_list:
        forecast, _ = nbeats_forward(window, blocks)
        if accumulated is None:
            horizon = len(forecast)
            accumulated = [0.0] * horizon
        for h in range(horizon):
            accumulated[h] += forecast[h]
    return [value / len(model_list) for value in accumulated]

"""The same architecture, with the basis as a type and the windows checked.

Two changes matter. The basis matrices are fixed constants determined by the
horizon and the basis choice, so they are built once at construction rather
than per forward pass. And the preconditions that silently produce nonsense -
a series shorter than the lookback, a polynomial degree high enough to
extrapolate absurdly, an ensemble quoted but not built - become explicit.
"""

from __future__ import annotations

import math
from dataclasses import dataclass
from enum import Enum
from typing import NamedTuple, Sequence


class BasisKind(Enum):
    """What family the block's output is constrained to lie in."""

    GENERIC = 'generic'
    TREND = 'trend'
    SEASONALITY = 'seasonality'


class InsufficientHistory(ValueError):
    """Raised when a series is shorter than the fixed lookback window.

    Its own type because the correct response is specific: fall back to a
    naive forecast for that series. The architecture cannot forecast it at
    all, and there is no partial-window path.
    """


class ShapeMismatch(ValueError):
    """Raised on a shape violation instead of broadcasting past it."""


class ImplausibleBasis(ValueError):
    """Raised for a basis that will extrapolate absurdly.

    A degree-6 polynomial fits history beautifully and reaches implausible
    values at the far end of the horizon, and nothing in the loss penalizes
    that. Refusing it is cheaper than discovering it in production.
    """


@dataclass(frozen=True)
class BasisSpec:
    """Frozen, because a basis is a structural prior, not a tunable."""

    kind: BasisKind
    degree: int = 2        # polynomial degree for TREND
    harmonics: int = 4     # harmonic count for SEASONALITY

    def __post_init__(self) -> None:
        if self.kind is BasisKind.TREND and not 1 <= self.degree <= 4:
            raise ImplausibleBasis(
                f'polynomial degree {self.degree} extrapolates unusably; keep it at 4 or below'
            )
        if self.kind is BasisKind.SEASONALITY and self.harmonics < 1:
            raise ShapeMismatch('at least one harmonic is required')

    def coefficient_count(self, length: int) -> int:
        """How many numbers the MLP head must emit for this basis."""
        if self.kind is BasisKind.GENERIC:
            return length
        if self.kind is BasisKind.TREND:
            return self.degree + 1
        return 2 * self.harmonics

    def build(self, length: int) -> list[list[float]]:
        """The basis matrix: length rows of coefficient_count.

        Built once and reused. It depends only on the length and the spec,
        so rebuilding it per forward pass is pure waste.
        """
        if length <= 0:
            raise ShapeMismatch('basis length must be positive')

        if self.kind is BasisKind.GENERIC:
            return [
                [1.0 if i == j else 0.0 for j in range(length)]
                for i in range(length)
            ]

        if self.kind is BasisKind.TREND:
            # Normalized time, so the coefficients mean the same thing
            # whatever the horizon is.
            return [
                [(index / length) ** power for power in range(self.degree + 1)]
                for index in range(length)
            ]

        rows: list[list[float]] = []
        for index in range(length):
            row: list[float] = []
            for harmonic in range(1, self.harmonics + 1):
                angle = 2.0 * math.pi * harmonic * index / length
                row.extend((math.cos(angle), math.sin(angle)))
            rows.append(row)
        return rows


@dataclass(frozen=True)
class StackConfig:
    """The shape contract, stated once and derived from rather than repeated."""

    lookback: int
    horizon: int
    hidden_width: int = 512
    layers: int = 4
    blocks: int = 3

    def __post_init__(self) -> None:
        if self.horizon < 1:
            raise ShapeMismatch('horizon must be at least one step')
        if self.lookback < self.horizon:
            raise ShapeMismatch('a lookback shorter than the horizon cannot work')

    @property
    def lookback_multiplier(self) -> float:
        """The strongest single lever, and the main source of ensemble
        diversity — so it is worth naming rather than leaving implicit."""
        return self.lookback / self.horizon


class BlockOutput(NamedTuple):
    """Both halves. The backcast is not a by-product — it is what makes the
    residual recursion work, and discarding it breaks the architecture."""

    backcast: list[float]
    forecast: list[float]


class Block:
    """One N-BEATS block: shared MLP trunk, two coefficient heads, two bases."""

    def __init__(
        self,
        config: StackConfig,
        backcast_spec: BasisSpec,
        forecast_spec: BasisSpec,
    ) -> None:
        self._config = config
        # Built once at construction: these are constants of the geometry.
        self._backcast_basis = backcast_spec.build(config.lookback)
        self._forecast_basis = forecast_spec.build(config.horizon)
        self._backcast_coefficients = backcast_spec.coefficient_count(config.lookback)
        self._forecast_coefficients = forecast_spec.coefficient_count(config.horizon)

        self._layers: list[tuple[list[list[float]], list[float]]] = []
        width = config.lookback
        for _ in range(config.layers):
            self._layers.append((
                [[0.0] * width for _ in range(config.hidden_width)],
                [0.0] * config.hidden_width,
            ))
            width = config.hidden_width

        self._backcast_head = [[0.0] * width for _ in range(self._backcast_coefficients)]
        self._forecast_head = [[0.0] * width for _ in range(self._forecast_coefficients)]

    def forward(self, window: Sequence[float]) -> BlockOutput:
        if len(window) != self._config.lookback:
            raise ShapeMismatch(
                f'expected a {self._config.lookback}-step window, got {len(window)}'
            )

        hidden: Sequence[float] = window
        for weight, bias in self._layers:
            hidden = [
                max(0.0, b + sum(w * v for w, v in zip(row, hidden)))
                for row, b in zip(weight, bias)
            ]

        # The heads predict COEFFICIENTS, not values. With a generic basis
        # that is vacuous; with trend or seasonality it is the whole point.
        theta_backcast = [sum(w * v for w, v in zip(row, hidden)) for row in self._backcast_head]
        theta_forecast = [sum(w * v for w, v in zip(row, hidden)) for row in self._forecast_head]

        return BlockOutput(
            backcast=[
                math.fsum(b * t for b, t in zip(row, theta_backcast))
                for row in self._backcast_basis
            ],
            forecast=[
                math.fsum(b * t for b, t in zip(row, theta_forecast))
                for row in self._forecast_basis
            ],
        )


class ResidualTrace(NamedTuple):
    """Forecast plus the residual norm after each block.

    The norms are returned deliberately: a block whose incoming residual is
    already near zero contributes parameters and inference time and learns
    nothing, and that is invisible unless it is measured.
    """

    forecast: list[float]
    residual_norms: list[float]


def forward(window: Sequence[float], blocks: Sequence[Block], horizon: int) -> ResidualTrace:
    """Doubly residual stacking: forecasts add, inputs shrink."""
    if not blocks:
        raise ShapeMismatch('a stack needs at least one block')

    residual = list(window)
    total = [0.0] * horizon
    norms: list[float] = []

    for block in blocks:
        output = block.forward(residual)
        residual = [r - b for r, b in zip(residual, output.backcast)]
        total = [t + f for t, f in zip(total, output.forecast)]
        norms.append(math.sqrt(math.fsum(value * value for value in residual)))

    return ResidualTrace(forecast=total, residual_norms=norms)


def mase(
    forecast: Sequence[float],
    actual: Sequence[float],
    history: Sequence[float],
    seasonal_period: int,
) -> float:
    """Mean absolute scaled error.

    Guard clauses first: a flat series gives a zero denominator, and a short
    one gives no denominator at all. Both are real and both would otherwise
    produce an infinity that poisons the batch average silently.
    """
    if len(forecast) != len(actual):
        raise ShapeMismatch('forecast and actual lengths disagree')
    if len(history) <= seasonal_period:
        raise InsufficientHistory('history is shorter than one seasonal period')

    differences = [
        abs(history[i] - history[i - seasonal_period])
        for i in range(seasonal_period, len(history))
    ]
    denominator = math.fsum(differences) / len(differences)
    if denominator == 0.0:
        raise ShapeMismatch('a constant history cannot scale the loss')

    numerator = math.fsum(abs(f - a) for f, a in zip(forecast, actual)) / len(forecast)
    return numerator / denominator


def window_or_fail(series: Sequence[float], config: StackConfig) -> list[float]:
    """The fixed lookback is structural. A series shorter than it is not
    forecastable by this architecture at all, and needs a naive fallback."""
    if len(series) < config.lookback:
        raise InsufficientHistory(
            f'{len(series)} observations is short of the {config.lookback}-step window'
        )
    return list(series[-config.lookback:])

"""Batched, and the whole stack is GEMMs. There is nothing sequential here.

The structural observation: N-BEATS has no recurrence and no attention, so
every layer is a dense matrix product and the only ordering constraint is
between blocks. That makes it the easiest forecaster in this reference to
make fast, and it is why an ensemble of thirty models is affordable at all.

Three changes:
  1. The MLP trunk, both coefficient heads and both basis projections become
     GEMMs over the whole batch. The bases are CONSTANT matrices, so they are
     built once and multiplied, never regenerated.
  2. The two heads share one GEMM: their outputs are concatenated columns of
     a single weight, so the trunk activations are read once rather than twice.
  3. The ensemble batches across models rather than looping. Members share a
     shape, so a (members, batch, lookback) tensor flows through as one
     batched contraction - which is what makes the ensemble cheap enough to
     be the method rather than a luxury.
"""

from __future__ import annotations

import numpy as np
from numpy.typing import NDArray

FLOAT = np.float32


def build_trend_basis(length: int, degree: int) -> NDArray[np.float32]:
    """(length, degree + 1) of normalized-time powers. A constant."""
    time = (np.arange(length, dtype=FLOAT) / FLOAT(length))[:, None]
    powers = np.arange(degree + 1, dtype=FLOAT)[None, :]
    return np.power(time, powers, dtype=FLOAT)


def build_seasonality_basis(length: int, harmonics: int) -> NDArray[np.float32]:
    """(length, 2 * harmonics) of interleaved cosines and sines. A constant."""
    time = np.arange(length, dtype=FLOAT)[:, None]
    frequency = np.arange(1, harmonics + 1, dtype=FLOAT)[None, :]
    angle = FLOAT(2.0 * np.pi) * frequency * time / FLOAT(length)

    basis = np.empty((length, 2 * harmonics), dtype=FLOAT)
    # Interleaved by strided assignment rather than by concatenate-and-
    # reorder, so no intermediate is materialized.
    basis[:, 0::2] = np.cos(angle)
    basis[:, 1::2] = np.sin(angle)
    return basis


class BatchedBlock:
    """One block, evaluated for a whole batch at once.

    The bases are stored as contiguous constants. A basis projection is then
    a GEMM against a fixed matrix — which is the cheapest possible way to
    impose the structural prior, and a good reason to prefer this form over
    generating the output values directly.
    """

    def __init__(
        self,
        trunk_weights: list[NDArray[np.float32]],
        trunk_biases: list[NDArray[np.float32]],
        head_weight: NDArray[np.float32],
        backcast_basis: NDArray[np.float32],
        forecast_basis: NDArray[np.float32],
    ) -> None:
        self._trunk_weights = [np.ascontiguousarray(w, dtype=FLOAT) for w in trunk_weights]
        self._trunk_biases = [np.ascontiguousarray(b, dtype=FLOAT) for b in trunk_biases]
        # ONE head weight whose columns are backcast coefficients followed by
        # forecast coefficients: the trunk output is read once, not twice.
        self._head_weight = np.ascontiguousarray(head_weight, dtype=FLOAT)
        self._backcast_basis = np.ascontiguousarray(backcast_basis, dtype=FLOAT)
        self._forecast_basis = np.ascontiguousarray(forecast_basis, dtype=FLOAT)
        self._split = backcast_basis.shape[1]

    def __call__(
        self, window: NDArray[np.float32]
    ) -> tuple[NDArray[np.float32], NDArray[np.float32]]:
        """window is (batch, lookback). Returns (backcast, forecast)."""
        hidden = window
        for weight, bias in zip(self._trunk_weights, self._trunk_biases):
            hidden = hidden @ weight
            hidden += bias
            # In-place ReLU: the trunk is the widest tensor in the block and
            # a temporary per layer is the largest avoidable allocation.
            np.maximum(hidden, 0.0, out=hidden)

        theta = hidden @ self._head_weight
        backcast = theta[:, :self._split] @ self._backcast_basis.T
        forecast = theta[:, self._split:] @ self._forecast_basis.T
        return backcast, forecast


def stack_forward(
    window: NDArray[np.float32],
    blocks: list[BatchedBlock],
) -> tuple[NDArray[np.float32], NDArray[np.float32]]:
    """Doubly residual stacking. Both recursions run in place.

    Returns the summed forecast and the per-block residual norms, which are
    the diagnostic for depth that has stopped paying.
    """
    if not blocks:
        raise ValueError('a stack needs at least one block')

    residual = window.copy()
    total: NDArray[np.float32] | None = None
    norms = np.empty((len(blocks), window.shape[0]), dtype=FLOAT)

    for index, block in enumerate(blocks):
        backcast, forecast = block(residual)
        # Both recursions in place: the residual shrinks and the forecast
        # accumulates without either allocating per block.
        residual -= backcast
        if total is None:
            total = forecast
        else:
            total += forecast
        norms[index] = np.linalg.norm(residual, axis=1)

    return total, norms


def ensemble_forward(
    windows: NDArray[np.float32],
    stacked_trunk: NDArray[np.float32],
    stacked_bias: NDArray[np.float32],
) -> NDArray[np.float32]:
    """One layer of every ensemble member at once.

    windows is (members, batch, lookback), stacked_trunk is
    (members, lookback, hidden). matmul broadcasts over the leading axis, so
    all members advance in a single batched GEMM rather than a Python loop -
    which is what turns a thirty-model ensemble from a luxury into the
    default configuration.
    """
    hidden = np.matmul(windows, stacked_trunk)
    hidden += stacked_bias[:, None, :]
    np.maximum(hidden, 0.0, out=hidden)
    return hidden


def mase(
    forecasts: NDArray[np.float32],
    actuals: NDArray[np.float32],
    histories: NDArray[np.float32],
    seasonal_period: int,
) -> NDArray[np.float32]:
    """Per-series MASE, vectorized over the batch.

    The seasonal-naive denominator is a single strided difference over the
    history axis — no loop, and no per-series Python at all.
    """
    if histories.shape[1] <= seasonal_period:
        raise ValueError('history is shorter than one seasonal period')

    seasonal_error = np.abs(histories[:, seasonal_period:] - histories[:, :-seasonal_period])
    denominator = seasonal_error.mean(axis=1, dtype=np.float32)
    # A constant series gives a zero denominator, which would produce an
    # infinity that silently poisons the batch average.
    if np.any(denominator == 0.0):
        raise ValueError('a constant history cannot scale the loss')

    numerator = np.abs(forecasts - actuals).mean(axis=1, dtype=np.float32)
    return numerator / denominator


def ensemble_spread(member_forecasts: NDArray[np.float32]) -> NDArray[np.float32]:
    """Disagreement across members, as a free uncertainty proxy.

    The model emits no interval, but a thirty-member ensemble already
    contains one. It is not calibrated and should not be presented as a
    predictive distribution — but it is a genuinely useful drift signal.
    """
    return member_forecasts.std(axis=0, dtype=np.float32)