Time-Series Diffusion & Foundation Models

"""Masked conditional diffusion for time series, transcribed literally.

The central idea is the MASK. A forecast and an imputation are the same
operation - some values are known and some are not:

    mask = 0  ->  observed, conditioning
    mask = 1  ->  unknown, predicted

Mask the tail and it is forecasting. Mask the middle and it is imputation. The
model does not distinguish, and randomizing the mask during training is what
makes one model handle both.

NOTATIONAL HAZARD, and a real source of bugs: `t` below is the DIFFUSION
timestep, not the position in the series. They index completely different axes
and are both conventionally called t.

Plain loops, no libraries.
"""

import math
import random

SEED = 37


def build_schedule(steps):
    """Cosine schedule, as in the DDPM entry. Fixed, not learned."""
    def f(index):
        angle = ((index / steps) + 0.008) / 1.008 * math.pi / 2.0
        return math.cos(angle) ** 2

    bars = [f(index) / f(0) for index in range(steps + 1)]
    alpha_bars = bars[1:]

    return {
        'alpha_bars': alpha_bars,
        'sqrt_alpha_bars': [math.sqrt(a) for a in alpha_bars],
        'sqrt_one_minus': [math.sqrt(1.0 - a) for a in alpha_bars],
    }


def scale_window(window, mask):
    """Per-window normalization, computed from the OBSERVED values only.

    This is what lets one model span series measured in units and series
    measured in millions. Two details matter. It uses only the observed
    positions, because including the masked ones would leak the answer. And
    it is part of the INFERENCE path rather than preprocessing - a scaler
    computed differently at serving time changes every forecast.
    """
    observed = [value for value, m in zip(window, mask) if m == 0]
    if not observed:
        return list(window), 0.0, 1.0

    mean = sum(observed) / len(observed)
    variance = sum((value - mean) ** 2 for value in observed) / len(observed)
    scale = math.sqrt(variance) or 1.0

    return [(value - mean) / scale for value in window], mean, scale


def random_mask(length, rng, strategy='mixed'):
    """Randomize the mask during training.

    A model trained only on tail masks forecasts and imputes badly; a model
    trained on mixed masks does both. The mask strategy determines which tasks
    the model can perform, which makes it a modelling decision rather than a
    data-loading detail.
    """
    if strategy == 'forecast' or (strategy == 'mixed' and rng.random() < 0.5):
        cut = rng.randrange(length // 2, length)
        return [0] * cut + [1] * (length - cut)

    # Block missingness: a contiguous gap somewhere in the middle.
    start = rng.randrange(0, length - length // 4)
    width = rng.randrange(1, length // 4 + 1)
    return [1 if start <= index < start + width else 0 for index in range(length)]


def denoiser(x_t, t, observed, mask, params):
    """Conditioned on the noisy target, the timestep, AND the observed values.

    The observed values and the mask are both inputs at every step. Supplying
    them only at initialization is the classic mistake - see reverse_step.
    """
    features = list(x_t) + list(observed) + [float(m) for m in mask] + [t / 100.0]

    hidden = []
    for row, bias in zip(params['w1'], params['b1']):
        total = bias + sum(w * v for w, v in zip(row, features))
        hidden.append(math.tanh(total))

    return [
        bias + sum(w * h for w, h in zip(row, hidden))
        for row, bias in zip(params['w2'], params['b2'])
    ]


def training_loss(window, params, schedule, rng):
    """Masked noise prediction. The loss applies ONLY where the mask is one.

    Everything outside the mask is conditioning, not a target - which is
    exactly what makes forecasting and imputation the same objective.
    """
    mask = random_mask(len(window), rng)
    scaled, _, _ = scale_window(window, mask)

    t = rng.randrange(len(schedule['alpha_bars']))
    signal = schedule['sqrt_alpha_bars'][t]
    noise_scale = schedule['sqrt_one_minus'][t]

    epsilon = [rng.gauss(0.0, 1.0) for _ in scaled]
    x_t = [signal * value + noise_scale * e for value, e in zip(scaled, epsilon)]

    # Observed positions are passed through unnoised as conditioning.
    observed = [0.0 if m else value for value, m in zip(scaled, mask)]
    predicted = denoiser(x_t, t, observed, mask, params)

    counted = sum(mask)
    if counted == 0:
        raise ValueError('a mask with nothing to predict trains on nothing')

    total = 0.0
    for e, p, m in zip(epsilon, predicted, mask):
        if m:
            total += (e - p) ** 2
    return total / counted


def reverse_step(x_t, t, observed, mask, params, schedule, rng, alphas):
    """One denoising step, with the observed values RE-IMPOSED.

    This is the step everyone forgets, and its absence is the single most
    common implementation error in this family.

    If the known positions are fixed once at initialization and then left
    alone, the denoiser drifts away from them as it works - it is filling a
    canvas and has no reason to respect a region it was not asked about. So
    the reconstruction ends up inconsistent with data you actually had.

    Re-noising the observed values to the CURRENT noise level and substituting
    them back at every step keeps the trajectory anchored to the real data.
    """
    predicted = denoiser(x_t, t, observed, mask, params)

    alpha = alphas[t]
    alpha_bar = schedule['alpha_bars'][t]
    coefficient = (1.0 - alpha) / math.sqrt(1.0 - alpha_bar)

    denoised = [
        (value - coefficient * p) / math.sqrt(alpha)
        for value, p in zip(x_t, predicted)
    ]

    if t == 0:
        return [d if m else o for d, o, m in zip(denoised, observed, mask)]

    # Re-noise the observed values to the PREVIOUS level and substitute back.
    previous = schedule['alpha_bars'][t - 1]
    signal = math.sqrt(previous)
    noise_scale = math.sqrt(1.0 - previous)

    return [
        d if m else signal * o + noise_scale * rng.gauss(0.0, 1.0)
        for d, o, m in zip(denoised, observed, mask)
    ]


def sample_trajectories(window, mask, params, schedule, alphas, count, rng):
    """Draw several completions. Their spread IS the predictive distribution.

    Cost note, because it decides feasibility: this is steps times count times
    however many series are in the panel. A hundred trajectories at fifty
    steps across ten thousand series is fifty million network evaluations for
    one forecast cycle - a different operational regime entirely from fitting
    a classical model per series.
    """
    scaled, mean, scale = scale_window(window, mask)
    observed = [0.0 if m else value for value, m in zip(scaled, mask)]

    trajectories = []
    for _ in range(count):
        x = [rng.gauss(0.0, 1.0) for _ in scaled]
        for t in range(len(schedule['alpha_bars']) - 1, -1, -1):
            x = reverse_step(x, t, observed, mask, params, schedule, rng, alphas)
        # Undo the per-window scaling. The scaler is part of inference.
        trajectories.append([value * scale + mean for value in x])

    return trajectories


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


def seasonal_naive_baseline(history, horizon, period):
    """The baseline that wins more often than the literature suggests.

    Worth having in the same file as the model: published comparisons for
    pretrained forecasters frequently omit properly tuned classical baselines,
    and this one is three lines.
    """
    return [history[-period + (index % period)] for index in range(horizon)]

"""The same model, with the mask as a type and the scaler carried with it.

Two things change. The mask becomes an explicit object that knows what task it
represents and validates that it asks for something, because it is the whole
modelling decision rather than a data-loading detail. And the per-window
scaler becomes part of the returned forecast, because it is part of the
INFERENCE path — a scaler computed differently at serving time changes every
forecast, with nothing to indicate it.
"""

from __future__ import annotations

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


class TaskKind(Enum):
    """What the mask asks for. Determined by its shape, not by configuration.

    Forecasting and imputation are the same operation with different masks, so
    naming the task is a property of the mask rather than a separate setting.
    """

    FORECAST = 'forecast'
    IMPUTATION = 'imputation'
    MIXED = 'mixed'


class EmptyMask(ValueError):
    """Raised when a mask asks for nothing.

    Its own type because the consequence is specific: the loss has no terms,
    so the batch contributes nothing and the training loop reports a perfectly
    healthy zero.
    """


class ContextOverflow(ValueError):
    """Raised when the history exceeds the pretrained context window.

    Its own type because the alternative is silent truncation: a pretrained
    model simply drops the excess and reports nothing, and the symptom is
    unexplained accuracy loss on exactly the series with the most history.
    """


class MissingScaler(ValueError):
    """Raised when a forecast is produced without its fitted scaler.

    Per-window scaling is part of the inference path, not preprocessing.
    Recomputing it differently at serving time changes every forecast.
    """


class ShapeMismatch(ValueError):
    """Raised on a dimension violation instead of computing nonsense."""


@dataclass(frozen=True)
class WindowScaler:
    """Per-window normalization, fitted on the OBSERVED values only.

    Two details matter and both are easy to get wrong. It must use only the
    observed positions, because including the masked ones leaks the answer
    into the scaling. And it travels with the forecast, because undoing it is
    part of producing the output rather than a separate step.
    """

    mean: float
    scale: float

    @staticmethod
    def fit(window: Sequence[float], mask: Sequence[int]) -> WindowScaler:
        observed = [value for value, m in zip(window, mask) if m == 0]
        if not observed:
            raise EmptyMask('cannot scale a window with no observed values')

        mean = math.fsum(observed) / len(observed)
        variance = math.fsum((value - mean) ** 2 for value in observed) / len(observed)
        # A constant window gives zero variance; scaling by one leaves it
        # unchanged rather than dividing by zero.
        return WindowScaler(mean=mean, scale=math.sqrt(variance) or 1.0)

    def apply(self, window: Sequence[float]) -> list[float]:
        return [(value - self.mean) / self.scale for value in window]

    def invert(self, window: Sequence[float]) -> list[float]:
        return [value * self.scale + self.mean for value in window]


class TargetMask:
    """Which positions are predicted. The whole modelling decision.

    A model trained only on tail masks forecasts and imputes badly; one
    trained on mixed masks does both. Making the mask a type with a task
    property means that choice is visible rather than buried in a data loader.
    """

    def __init__(self, values: Sequence[int]) -> None:
        if not any(values):
            raise EmptyMask(
                'the mask asks for nothing; the loss will have no terms and the '
                'batch will report a healthy zero'
            )
        if all(values):
            raise EmptyMask('the mask leaves nothing observed; there is no conditioning')
        self._values = list(values)

    def __len__(self) -> int:
        return len(self._values)

    def __iter__(self):
        return iter(self._values)

    @property
    def predicted_count(self) -> int:
        return sum(self._values)

    @property
    def task(self) -> TaskKind:
        """Derived from the shape rather than declared.

        A mask that is zeros then ones is a forecast; anything else with an
        interior gap is imputation. Deriving it means a mislabelled mask is
        impossible rather than merely unlikely.
        """
        first_one = next(index for index, value in enumerate(self._values) if value)
        if all(self._values[first_one:]):
            return TaskKind.FORECAST
        return TaskKind.IMPUTATION

    @staticmethod
    def forecast(length: int, horizon: int) -> TargetMask:
        if horizon >= length:
            raise EmptyMask('a horizon covering the whole window leaves no context')
        return TargetMask([0] * (length - horizon) + [1] * horizon)

    @staticmethod
    def random(length: int, rng: random.Random) -> TargetMask:
        """Randomized during training, so one model handles both tasks."""
        if rng.random() < 0.5:
            return TargetMask.forecast(length, rng.randrange(1, length // 2))
        start = rng.randrange(0, length - length // 4)
        width = rng.randrange(1, length // 4 + 1)
        return TargetMask(
            [1 if start <= index < start + width else 0 for index in range(length)]
        )


@dataclass(frozen=True)
class ContextWindow:
    """The pretrained window, which is a hard limit rather than a default."""

    length: int

    def assert_fits(self, history: Sequence[float]) -> None:
        """Guard clause for the failure that is otherwise silent.

        A pretrained model given a longer history simply truncates and reports
        nothing, so the symptom is unexplained accuracy loss on exactly the
        series with the most data.
        """
        if len(history) > self.length:
            raise ContextOverflow(
                f'{len(history)} observations exceed the {self.length}-step pretrained '
                'context; the excess would be silently truncated'
            )


class Forecast(NamedTuple):
    """Trajectories, the scaler, and the settings that produced them.

    The scaler travels with the forecast because undoing it is part of
    producing the output. The step and trajectory counts travel because
    quality and cost are on a dial — a forecast reported without them is not
    comparable to any other.
    """

    trajectories: list[list[float]]
    scaler: WindowScaler
    steps: int
    trajectory_count: int

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

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

        The joint-distribution advantage: summing per-step quantiles instead
        destroys the correlation between steps and is wrong in the unsafe
        direction for any aggregating decision.
        """
        totals = sorted(math.fsum(path) for path in self.trajectories)
        return totals[min(int(level * len(totals)), len(totals) - 1)]

    def tail_is_estimable(self, level: float, minimum_order_statistics: int = 5) -> bool:
        """Whether there are enough trajectories for this quantile to mean
        anything. A 1st percentile from fifty trajectories is one order
        statistic and is pure noise."""
        tail = min(level, 1.0 - level)
        return tail * self.trajectory_count >= minimum_order_statistics


def masked_loss(
    epsilon: Sequence[float], predicted: Sequence[float], mask: TargetMask
) -> float:
    """The loss applies ONLY where the mask is one.

    Everything outside it is conditioning rather than a target, which is
    exactly what makes forecasting and imputation the same objective.
    """
    if len(epsilon) != len(predicted) or len(epsilon) != len(mask):
        raise ShapeMismatch('noise, prediction and mask widths disagree')

    total = math.fsum(
        (e - p) ** 2 for e, p, m in zip(epsilon, predicted, mask) if m
    )
    return total / mask.predicted_count


def reimpose_observed(
    denoised: Sequence[float],
    observed: Sequence[float],
    mask: TargetMask,
    alpha_bar_previous: float,
    rng: random.Random,
) -> list[float]:
    """Re-noise the observed values and substitute them back. EVERY step.

    The single most commonly omitted step in this family. If the known
    positions are fixed once at initialization and then left alone, the
    denoiser drifts away from them as it works — it is filling a canvas and
    has no reason to respect a region it was not asked about, so the
    reconstruction ends up inconsistent with data you actually had.
    """
    signal = math.sqrt(alpha_bar_previous)
    noise_scale = math.sqrt(1.0 - alpha_bar_previous)

    return [
        d if m else signal * o + noise_scale * rng.gauss(0.0, 1.0)
        for d, o, m in zip(denoised, observed, mask)
    ]


def inference_evaluations(steps: int, trajectories: int, series: int) -> int:
    """The arithmetic that decides feasibility, computed rather than assumed.

    Steps times trajectories times series. A hundred trajectories at fifty
    steps across ten thousand series is fifty million network evaluations for
    one forecast cycle — a different operational regime entirely from fitting
    a classical model per series, and routinely under-planned by an order of
    magnitude.
    """
    return steps * trajectories * series

"""Batched over series AND trajectories, which is the only lever that helps.

The structural situation is worth stating precisely, because it differs from
every other diffusion entry. Sampling cost here is STEPS x TRAJECTORIES x
SERIES, and only one of those three is sequential. Steps must happen in order;
trajectories and series are entirely independent and fold into one batch axis.

So the optimization is: collapse trajectories and series into a single leading
dimension, advance the whole thing through each sequential step together, and
accept that the step count is the only part that cannot be parallelized away.

Three changes:
  1. (series x trajectories) becomes one batch axis.
  2. The mask and the observed values are precomputed once, not per step.
  3. Re-imposition is a fused select rather than a branch per element.
"""

from __future__ import annotations

import numpy as np
from numpy.typing import NDArray

FLOAT = np.float32


def fit_window_scaler(
    windows: NDArray[np.float32], mask: NDArray[np.bool_]
) -> tuple[NDArray[np.float32], NDArray[np.float32]]:
    """Per-window mean and scale from the OBSERVED positions only.

    Using the masked positions would leak the answer into the scaling, so the
    reduction is weighted by the complement of the mask rather than taken over
    the whole window — one pass, no boolean indexing, no compacted copy.
    """
    observed = ~mask
    counts = observed.sum(axis=1, keepdims=True).astype(FLOAT)
    counts = np.maximum(counts, 1.0)

    mean = (windows * observed).sum(axis=1, keepdims=True, dtype=FLOAT) / counts
    centred = (windows - mean) * observed
    variance = np.einsum('ij,ij->i', centred, centred, optimize=True)[:, None] / counts
    # A constant window gives zero variance; one leaves it unchanged rather
    # than dividing by zero.
    scale = np.maximum(np.sqrt(variance), 1e-6)
    return mean, scale


def expand_for_trajectories(
    windows: NDArray[np.float32], mask: NDArray[np.bool_], trajectories: int
) -> tuple[NDArray[np.float32], NDArray[np.bool_]]:
    """Fold (series, trajectories) into ONE batch axis.

    This is the whole optimization. Trajectories and series are both entirely
    independent, so they belong in the same dimension; only the diffusion
    steps are sequential, and nothing rearranges that.

    repeat rather than tile, so each series' trajectories are contiguous —
    which keeps the per-series quantile reduction at the end a contiguous
    read rather than a strided one.
    """
    return (
        np.repeat(windows, trajectories, axis=0),
        np.repeat(mask, trajectories, axis=0),
    )


class BatchedSampler:
    """Every trajectory of every series advances through each step together."""

    def __init__(self, alpha_bars: NDArray[np.float32]) -> None:
        self._alpha_bars = np.ascontiguousarray(alpha_bars, dtype=FLOAT)
        alphas = np.empty_like(self._alpha_bars)
        alphas[0] = self._alpha_bars[0]
        alphas[1:] = self._alpha_bars[1:] / self._alpha_bars[:-1]

        # Every coefficient precomputed: these are constants of the schedule
        # and recomputing them inside the step loop is pure waste.
        self._sqrt_alpha_bars = np.sqrt(self._alpha_bars)
        self._sqrt_one_minus = np.sqrt(1.0 - self._alpha_bars)
        self._reverse_coefficient = (1.0 - alphas) / self._sqrt_one_minus
        self._root_alphas = np.sqrt(alphas)

    @property
    def steps(self) -> int:
        return len(self._alpha_bars)

    def reverse_step(
        self,
        x_t: NDArray[np.float32],
        predicted: NDArray[np.float32],
        observed: NDArray[np.float32],
        mask: NDArray[np.bool_],
        t: int,
        rng: np.random.Generator,
    ) -> NDArray[np.float32]:
        """One denoising step with the observed values RE-IMPOSED.

        The step everyone forgets, and its absence is the most common
        implementation error in this family: if the known positions are fixed
        once at initialization the denoiser drifts away from them, because it
        is filling a canvas and has no reason to respect a region it was not
        asked about.

        Written as a fused select rather than a branch per element — the mask
        is a boolean array, so the choice is a where over the whole batch.
        """
        # In place over the prediction buffer: the denoised state replaces it.
        out = predicted
        out *= -self._reverse_coefficient[t]
        out += x_t
        out /= self._root_alphas[t]

        if t == 0:
            return np.where(mask, out, observed)

        # Re-noise the observed values to the PREVIOUS level, then select.
        previous = self._alpha_bars[t - 1]
        noise = rng.standard_normal(observed.shape, dtype=FLOAT)
        noise *= FLOAT(np.sqrt(1.0 - previous))
        noise += FLOAT(np.sqrt(previous)) * observed

        return np.where(mask, out, noise)


def masked_loss(
    epsilon: NDArray[np.float32],
    predicted: NDArray[np.float32],
    mask: NDArray[np.bool_],
) -> np.float32:
    """The loss applies ONLY where the mask is one.

    The mask becomes a multiplicative weight rather than a boolean index,
    which avoids allocating a compacted copy of every predicted position —
    and the count comes from the same array.
    """
    counted = mask.sum(dtype=np.int64)
    if counted == 0:
        raise ValueError('the mask asks for nothing; the loss has no terms')

    residual = predicted
    residual -= epsilon
    residual *= mask
    return np.einsum('ij,ij->', residual, residual, optimize=True) / counted


def trajectory_quantiles(
    trajectories: NDArray[np.float32],
    series: int,
    levels: NDArray[np.float32],
) -> NDArray[np.float32]:
    """(series, levels, horizon) from a flat (series * trajectories, horizon).

    The reshape is free because expand_for_trajectories used repeat rather
    than tile, so each series' trajectories are already contiguous.
    """
    reshaped = trajectories.reshape(series, -1, trajectories.shape[1])
    return np.quantile(reshaped, levels, axis=1).transpose(1, 0, 2)


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

    The joint-distribution advantage: summing per-step quantiles instead
    destroys the correlation between steps and is wrong in the unsafe
    direction for any decision that aggregates over the horizon.
    """
    reshaped = trajectories.reshape(series, -1, trajectories.shape[1])
    totals = reshaped.sum(axis=2, dtype=FLOAT)
    return np.quantile(totals, levels, axis=1).T


def seasonal_naive(
    histories: NDArray[np.float32], horizon: int, period: int
) -> NDArray[np.float32]:
    """The baseline that wins more often than the literature suggests.

    One strided gather for a whole panel. Worth having in the same file as the
    model: published comparisons for pretrained forecasters frequently omit
    properly tuned classical baselines, and this costs nothing to run.
    """
    positions = -period + (np.arange(horizon) % period)
    return histories[:, positions]


def inference_evaluations(steps: int, trajectories: int, series: int) -> int:
    """The arithmetic that decides feasibility, computed rather than assumed.

    Only ONE of these three factors is sequential. Trajectories and series
    both fold into the batch, so the wall clock is set by steps and the
    hardware — but the total work is the product, and it is routinely
    under-planned by an order of magnitude.
    """
    return steps * trajectories * series