DeepAR

"""DeepAR, transcribed the way the paper reads.

The model emits DISTRIBUTION PARAMETERS, not values:

    h_t    = LSTM(h_{t-1}, y_{t-1} / nu, x_t)      recurrent state
    mu_t   = nu * (w_mu . h_t)                      location
    sigma_t= nu * softplus(w_sigma . h_t)           scale, forced positive
    loss   = -log N(y_t | mu_t, sigma_t)            maximum likelihood

nu is the series' own mean. Dividing the input by it and multiplying the
outputs back is the unglamorous detail that lets one model span six orders of
magnitude - without it the large series dominate the gradient entirely.

Plain loops, one LSTM cell written out, Gaussian likelihood.
"""

import math
import random

SEED = 7


def sigmoid(value):
    if value >= 0.0:
        return 1.0 / (1.0 + math.exp(-value))
    positive = math.exp(value)
    return positive / (1.0 + positive)


def softplus(value):
    """log(1 + e^x), which is positive everywhere.

    The scale parameter of a distribution MUST be positive, and that is a
    hard constraint the architecture has to enforce rather than hope for.
    """
    if value > 30.0:
        return value  # e^x overflows and log1p(e^x) is x to full precision
    return math.log1p(math.exp(value))


def matvec(weight, vector, bias):
    return [
        bias[i] + sum(w * v for w, v in zip(row, vector))
        for i, row in enumerate(weight)
    ]


def lstm_cell(inputs, hidden, cell, params):
    """One LSTM step. Four gates over the concatenated input and state."""
    joined = list(inputs) + list(hidden)

    forget = [sigmoid(v) for v in matvec(params['wf'], joined, params['bf'])]
    write = [sigmoid(v) for v in matvec(params['wi'], joined, params['bi'])]
    candidate = [math.tanh(v) for v in matvec(params['wc'], joined, params['bc'])]
    expose = [sigmoid(v) for v in matvec(params['wo'], joined, params['bo'])]

    new_cell = [f * c + i * g for f, c, i, g in zip(forget, cell, write, candidate)]
    new_hidden = [o * math.tanh(c) for o, c in zip(expose, new_cell)]
    return new_hidden, new_cell


def series_scale(history):
    """nu = 1 + mean(history). The plus one keeps an all-zero series usable."""
    return 1.0 + sum(history) / len(history)


def gaussian_log_likelihood(value, mu, sigma):
    """log N(value | mu, sigma). This IS the loss, negated."""
    z = (value - mu) / sigma
    return -0.5 * z * z - math.log(sigma) - 0.5 * math.log(2.0 * math.pi)


def negative_binomial_log_likelihood(count, mu, alpha):
    """The likelihood that is actually right for demand data.

    Gaussian on counts puts probability mass on negative demand, which makes
    every lower quantile meaningless. Nothing in the training loss objects to
    this - it simply produces a confidently incoherent interval.
    """
    shape = 1.0 / alpha
    return (
        math.lgamma(count + shape)
        - math.lgamma(shape)
        - math.lgamma(count + 1.0)
        + shape * math.log(shape / (shape + mu))
        + count * math.log(mu / (shape + mu))
    )


def train_sequence(history, covariates, params, scale):
    """Teacher forcing: the recurrence is unrolled against TRUE observations.

    That is what makes training one parallel graph - and it is also the
    source of exposure bias, because inference conditions on the model's own
    samples instead and the recurrence has never seen that state.
    """
    hidden = [0.0] * params['hidden_size']
    cell = [0.0] * params['hidden_size']
    total_log_likelihood = 0.0

    for t in range(1, len(history)):
        previous = history[t - 1] / scale
        inputs = [previous] + list(covariates[t])
        hidden, cell = lstm_cell(inputs, hidden, cell, params)

        mu = scale * sum(w * h for w, h in zip(params['w_mu'], hidden))
        sigma = scale * softplus(sum(w * h for w, h in zip(params['w_sigma'], hidden)))
        total_log_likelihood += gaussian_log_likelihood(history[t], mu, sigma)

    return -total_log_likelihood, hidden, cell


def sample_path(hidden, cell, last_value, future_covariates, params, scale, rng):
    """One trajectory: draw, feed back, repeat.

    This is the whole reason DeepAR gives a JOINT distribution over the
    horizon rather than independent per-step quantiles. Each step's sample
    genuinely conditions the next, so uncertainty compounds correctly.
    """
    path = []
    previous = last_value

    for t in range(len(future_covariates)):
        inputs = [previous / scale] + list(future_covariates[t])
        hidden, cell = lstm_cell(inputs, hidden, cell, params)

        mu = scale * sum(w * h for w, h in zip(params['w_mu'], hidden))
        sigma = scale * softplus(sum(w * h for w, h in zip(params['w_sigma'], hidden)))

        drawn = rng.gauss(mu, sigma)
        path.append(drawn)
        previous = drawn  # the sample becomes the next input

    return path


def forecast(history, covariates, future_covariates, params, num_paths=200):
    """Warm the state on history, then draw many trajectories."""
    rng = random.Random(SEED)
    scale = series_scale(history)

    _, hidden, cell = train_sequence(history, covariates, params, scale)

    return [
        sample_path(hidden, cell, history[-1], future_covariates, params, scale, rng)
        for _ in range(num_paths)
    ]


def quantile_from_paths(paths, horizon, level):
    """Any quantile is just a percentile of the sampled trajectories."""
    values = sorted(path[horizon] for path in paths)
    index = min(int(level * len(values)), len(values) - 1)
    return values[index]


def lead_time_quantile(paths, level):
    """The joint advantage, made concrete.

    Total demand over a lead time is NOT the sum of per-step medians. Summing
    WITHIN each path first and taking the quantile afterwards preserves the
    correlation between steps; the other order destroys it, and the error is
    usually in the unsafe direction.
    """
    totals = sorted(sum(path) for path in paths)
    index = min(int(level * len(totals)), len(totals) - 1)
    return totals[index]

"""The same model, with the likelihood as a type and the scale as a contract.

Two things change. The likelihood family becomes a protocol with its own
parameter constraints and sampler, because choosing it is a modelling
decision rather than a flag - and a Gaussian head on count data produces
negative quantiles that nothing in the training loss objects to. And the
per-series scale becomes an object carried with the model, because
recomputing it at inference silently changes every forecast.
"""

from __future__ import annotations

import math
import random
from dataclasses import dataclass
from typing import NamedTuple, Protocol, Sequence


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


class LikelihoodMisspecified(ValueError):
    """Raised when the data cannot have come from the chosen family.

    Its own type because the failure is silent otherwise: a Gaussian head on
    count data trains happily and emits negative lower quantiles, and no
    metric in the training loop will mention it.
    """


class UnscaledSeries(ValueError):
    """Raised when a forecast is attempted without the fitted scale.

    The scale is part of the model artifact. Recomputing it at inference from
    a different window changes every forecast, with nothing logged.
    """


def softplus(value: float) -> float:
    """log(1 + e^x), positive everywhere, with the overflow branch.

    Above about 30 the exponential overflows while log1p(exp(x)) equals x to
    full double precision, so the branch is correctness rather than speed.
    """
    if value > 30.0:
        return value
    return math.log1p(math.exp(value))


class Likelihood(Protocol):
    """A distribution family the network emits parameters for.

    Made a protocol rather than an enum because each family owns three
    genuinely different things: how many parameters it needs, how to
    constrain them, and how to sample from it.
    """

    @property
    def parameter_count(self) -> int: ...

    def constrain(self, raw: Sequence[float], scale: float) -> tuple[float, ...]: ...

    def log_probability(self, value: float, parameters: tuple[float, ...]) -> float: ...

    def sample(self, parameters: tuple[float, ...], rng: random.Random) -> float: ...

    def validate_data(self, values: Sequence[float]) -> None: ...


@dataclass(frozen=True)
class Gaussian:
    """Correct for genuinely continuous quantities, and only those."""

    @property
    def parameter_count(self) -> int:
        return 2

    def constrain(self, raw: Sequence[float], scale: float) -> tuple[float, ...]:
        # Both parameters are multiplied back by the series scale, so the
        # network only ever works in normalized space.
        return scale * raw[0], scale * softplus(raw[1])

    def log_probability(self, value: float, parameters: tuple[float, ...]) -> float:
        mu, sigma = parameters
        z = (value - mu) / sigma
        return -0.5 * z * z - math.log(sigma) - 0.5 * math.log(2.0 * math.pi)

    def sample(self, parameters: tuple[float, ...], rng: random.Random) -> float:
        mu, sigma = parameters
        return rng.gauss(mu, sigma)

    def validate_data(self, values: Sequence[float]) -> None:
        """Guard clause for the failure with no other symptom."""
        if all(float(value).is_integer() and value >= 0.0 for value in values):
            raise LikelihoodMisspecified(
                'every observation is a non-negative integer; a Gaussian head will '
                'emit negative lower quantiles. Use NegativeBinomial.'
            )


@dataclass(frozen=True)
class NegativeBinomial:
    """The right choice for count and intermittent demand data."""

    @property
    def parameter_count(self) -> int:
        return 2

    def constrain(self, raw: Sequence[float], scale: float) -> tuple[float, ...]:
        # Mean scales with the series; dispersion does NOT, since it is a
        # shape parameter and scaling it would change the distribution family
        # rather than its location.
        return scale * softplus(raw[0]), softplus(raw[1])

    def log_probability(self, value: float, parameters: tuple[float, ...]) -> float:
        mu, alpha = parameters
        shape = 1.0 / alpha
        return (
            math.lgamma(value + shape)
            - math.lgamma(shape)
            - math.lgamma(value + 1.0)
            + shape * math.log(shape / (shape + mu))
            + value * math.log(mu / (shape + mu))
        )

    def sample(self, parameters: tuple[float, ...], rng: random.Random) -> float:
        # Gamma-Poisson mixture: the standard construction, and the reason
        # the family handles overdispersion that Poisson cannot.
        mu, alpha = parameters
        shape = 1.0 / alpha
        rate = rng.gammavariate(shape, mu / shape)
        return float(_poisson(rate, rng))

    def validate_data(self, values: Sequence[float]) -> None:
        if any(value < 0.0 for value in values):
            raise LikelihoodMisspecified('a count likelihood cannot model negative values')


def _poisson(rate: float, rng: random.Random) -> int:
    """Knuth's method. Fine for the small rates demand data produces."""
    threshold = math.exp(-rate)
    count, product = 0, 1.0
    while product > threshold:
        product *= rng.random()
        count += 1
        if count > 10_000:  # bounded: a huge rate would otherwise spin forever
            break
    return count - 1


@dataclass(frozen=True)
class SeriesScale:
    """The per-series normalizer, carried WITH the model.

    Frozen and explicit because recomputing it at inference from a different
    window silently changes every forecast, and that is one of the most
    common production errors with this model.
    """

    value: float

    @staticmethod
    def from_history(history: Sequence[float]) -> SeriesScale:
        if not history:
            raise ShapeMismatch('cannot scale an empty series')
        # The plus one keeps an all-zero series usable rather than dividing
        # by zero, which real catalogues contain in quantity.
        return SeriesScale(1.0 + math.fsum(history) / len(history))


class ForecastPaths(NamedTuple):
    """Trajectories, not quantiles.

    The paths come back rather than summary statistics because the joint
    structure is the whole point: a quantile of a SUM over the horizon cannot
    be recovered from per-step quantiles.
    """

    paths: list[list[float]]
    scale: SeriesScale
    seed: int

    def marginal_quantile(self, horizon: int, level: float) -> float:
        values = sorted(path[horizon] for path in self.paths)
        return values[min(int(level * len(values)), len(values) - 1)]

    def aggregate_quantile(self, level: float) -> float:
        """Sum WITHIN each path, then take the quantile.

        This ordering is the joint-distribution advantage. Reversing it -
        summing per-step quantiles - destroys the correlation between steps
        and is wrong in the unsafe direction.
        """
        totals = sorted(math.fsum(path) for path in self.paths)
        return totals[min(int(level * len(totals)), len(totals) - 1)]


def forecast(
    history: Sequence[float],
    future_covariates: Sequence[Sequence[float]],
    step: object,
    likelihood: Likelihood,
    scale: SeriesScale | None = None,
    num_paths: int = 200,
    seed: int = 7,
) -> ForecastPaths:
    """Warm the state on history, then draw trajectories.

    `scale` must be the one fitted with the model. Passing None is refused
    rather than silently recomputed, because a recomputed scale is a
    different model.
    """
    if scale is None:
        raise UnscaledSeries(
            'the fitted per-series scale must be supplied; recomputing it here '
            'would silently change every forecast'
        )
    if num_paths < 50:
        raise ShapeMismatch(
            f'{num_paths} paths cannot estimate a tail quantile; use at least 50'
        )

    likelihood.validate_data(history)
    rng = random.Random(seed)
    paths: list[list[float]] = []

    for _ in range(num_paths):
        state = getattr(step, 'warm')(history, scale.value)
        previous = history[-1]
        path: list[float] = []
        for covariate_row in future_covariates:
            state, raw = getattr(step, 'advance')(state, previous / scale.value, covariate_row)
            parameters = likelihood.constrain(raw, scale.value)
            drawn = likelihood.sample(parameters, rng)
            path.append(drawn)
            previous = drawn  # the sample becomes the next input
        paths.append(path)

    return ForecastPaths(paths=paths, scale=scale, seed=seed)


def empirical_coverage(
    paths_per_series: Sequence[ForecastPaths],
    actuals: Sequence[Sequence[float]],
    level: float,
) -> float:
    """What fraction of actuals fall below the nominal quantile.

    Calibration is a separate question from accuracy and it is the one that
    matters operationally: an over-confident model can have excellent median
    error and still cause systematic stockouts.
    """
    hits = 0
    total = 0
    for forecast_paths, actual in zip(paths_per_series, actuals):
        for horizon, observed in enumerate(actual):
            if observed <= forecast_paths.marginal_quantile(horizon, level):
                hits += 1
            total += 1
    if total == 0:
        raise ShapeMismatch('no observations to measure coverage against')
    return hits / total

"""Sampling is the cost, so sampling is what batches.

The structural observation: training is teacher-forced and therefore one
parallel graph, but inference rolls S paths forward step by step and that is
two orders of magnitude more work. So the paths become a BATCH DIMENSION -
all S trajectories take step t together, and the Python loop is over the
horizon (tens) rather than over the paths (hundreds).

Three changes:
  1. The four LSTM gates share ONE weight with 4H columns, so the joined
     input is read once rather than four times.
  2. Sample paths batch. (paths, hidden) advances as one GEMM per step.
  3. Series batch on top of that, so the whole panel and every path advance
     in a single (series * paths, hidden) product per step.
"""

from __future__ import annotations

import numpy as np
from numpy.typing import NDArray

FLOAT = np.float32


def softplus_(values: NDArray[np.float32]) -> NDArray[np.float32]:
    """In-place, with the overflow branch as a select rather than a loop.

    Above about 30 the exponential overflows while log1p(exp(x)) equals x to
    full precision, so the large branch is correctness and not speed.
    """
    large = values > 30.0
    safe = np.where(large, 0.0, values)
    np.exp(safe, out=safe)
    np.log1p(safe, out=safe)
    return np.where(large, values, safe)


class BatchedLstm:
    """One cell, evaluated for an arbitrary leading batch.

    The four gates share a single weight of 4H columns. Four separate
    products would read the joined input four times; one reads it once, and
    the split afterwards is a view rather than a copy.
    """

    def __init__(self, input_size: int, hidden_size: int) -> None:
        self._hidden = hidden_size
        self._weight = np.zeros((input_size + hidden_size, 4 * hidden_size), dtype=FLOAT)
        self._bias = np.zeros(4 * hidden_size, dtype=FLOAT)

    def step(
        self,
        inputs: NDArray[np.float32],
        hidden: NDArray[np.float32],
        cell: NDArray[np.float32],
    ) -> tuple[NDArray[np.float32], NDArray[np.float32]]:
        """inputs (batch, input_size), hidden/cell (batch, hidden_size)."""
        joined = np.concatenate((inputs, hidden), axis=1)
        gates = joined @ self._weight
        gates += self._bias

        h = self._hidden
        forget, write, candidate, expose = (
            gates[:, :h], gates[:, h:2 * h], gates[:, 2 * h:3 * h], gates[:, 3 * h:],
        )

        # Three sigmoids and a tanh, all in place over the gates buffer, so
        # the activated gates never exist as separate arrays.
        for block in (forget, write, expose):
            np.negative(block, out=block)
            np.exp(block, out=block)
            block += 1.0
            np.reciprocal(block, out=block)
        np.tanh(candidate, out=candidate)

        cell *= forget
        cell += write * candidate
        new_hidden = np.tanh(cell) * expose
        return new_hidden, cell


def sample_paths(
    lstm: BatchedLstm,
    warm_hidden: NDArray[np.float32],
    warm_cell: NDArray[np.float32],
    last_values: NDArray[np.float32],
    future_covariates: NDArray[np.float32],
    mu_weight: NDArray[np.float32],
    sigma_weight: NDArray[np.float32],
    scale: NDArray[np.float32],
    num_paths: int,
    rng: np.random.Generator,
) -> NDArray[np.float32]:
    """Returns (series, paths, horizon).

    Every path of every series advances together. The Python loop is over the
    HORIZON - tens of iterations - rather than over hundreds of paths times
    thousands of series, which is the whole structural change.
    """
    num_series, hidden_size = warm_hidden.shape
    horizon, covariate_width = future_covariates.shape[1], future_covariates.shape[2]
    rows = num_series * num_paths

    # Tile the warmed state across paths. Every path of a series starts from
    # the same state and diverges only through its own draws.
    hidden = np.repeat(warm_hidden, num_paths, axis=0).astype(FLOAT)
    cell = np.repeat(warm_cell, num_paths, axis=0).astype(FLOAT)
    previous = np.repeat(last_values, num_paths).astype(FLOAT)
    row_scale = np.repeat(scale, num_paths).astype(FLOAT)

    out = np.empty((rows, horizon), dtype=FLOAT)
    inputs = np.empty((rows, 1 + covariate_width), dtype=FLOAT)

    for step in range(horizon):
        # Pre-allocated input buffer: the concatenation of the fed-back value
        # and the covariates is written in place every step.
        np.divide(previous, row_scale, out=inputs[:, 0])
        inputs[:, 1:] = np.repeat(future_covariates[:, step, :], num_paths, axis=0)

        hidden, cell = lstm.step(inputs, hidden, cell)

        mu = (hidden @ mu_weight) * row_scale
        sigma = softplus_(hidden @ sigma_weight) * row_scale

        # One vectorized draw for every path of every series at once. This is
        # the operation the whole restructuring exists to batch.
        drawn = rng.normal(mu, sigma).astype(FLOAT)
        out[:, step] = drawn
        previous = drawn  # the sample becomes the next input

    return out.reshape(num_series, num_paths, horizon)


def marginal_quantiles(
    paths: NDArray[np.float32], levels: NDArray[np.float32]
) -> NDArray[np.float32]:
    """(series, levels, horizon) from (series, paths, horizon).

    One partition-based quantile over the path axis, not a sort: the quantile
    needs a selection, and ordering the rest is discarded work.
    """
    return np.quantile(paths, levels, axis=1, method='linear').transpose(1, 0, 2)


def aggregate_quantiles(
    paths: NDArray[np.float32], levels: NDArray[np.float32]
) -> NDArray[np.float32]:
    """Sum WITHIN each path, then take the quantile over paths.

    This ordering is the joint-distribution advantage, and reversing it -
    summing per-step quantiles - destroys the correlation between steps. The
    difference is large for any aggregating decision and usually unsafe.
    """
    totals = paths.sum(axis=2, dtype=FLOAT)
    return np.quantile(totals, levels, axis=1).T


def empirical_coverage(
    paths: NDArray[np.float32],
    actuals: NDArray[np.float32],
    level: float,
) -> float:
    """Fraction of actuals below the nominal quantile, in one pass.

    Calibration is separate from accuracy and it is the operational one: an
    over-confident model can have excellent median error and still cause
    systematic stockouts.
    """
    threshold = np.quantile(paths, level, axis=1)
    return float((actuals <= threshold).mean())


def path_count_for_tail(level: float, relative_error: float = 0.1) -> int:
    """How many paths a tail quantile actually needs.

    The standard error of an empirical quantile scales as
    sqrt(p(1-p)/n) / density, so a far tail needs far more paths than the
    body. Worth computing rather than defaulting to 100 and hoping - a 1st
    percentile from 100 paths is one order statistic and is pure noise.
    """
    tail = min(level, 1.0 - level)
    return int(np.ceil(tail * (1.0 - tail) / (relative_error * tail) ** 2))