src/adapters/baseline.py

"""Naive baselines, deliberately first-class citizens of the Arena.

A foundation model that cannot out-forecast "tomorrow looks like today" is not
worth the GPU time, and a leaderboard that quietly omits that comparison is
flattering rather than useful. So the baselines are enrolled, archived and
graded on exactly the same footing as everything else, and they appear on the
standings table next to the models they embarrass.

They are also the only family that runs with no weights, no network and no
torch, which makes them the fixture the test suite reasons about: capability
gating, determinism and coverage arithmetic are all checked here first, where
the expected answer can be worked out by hand.

`random-walk` is the reference. Its band comes from the empirical distribution
of historical log returns, scaled by sqrt(h) -- which is the correct widening
for a driftless random walk and is why a well-calibrated foundation model
should be able to beat it on error but will struggle to beat it on coverage.
"""

from __future__ import annotations

import numpy as np
import pandas as pd

from .. import config
from .base import (OHLCV_COLUMNS, OUTPUT_OHLCV_PATHS, OUTPUT_QUANTILE_LINE,
                   AdapterError, Capabilities, ForecastAdapter,
                   ForecastResult, check_context)

# `kind` -> how the median path is built, and what the adapter emits.
BASELINE_MODELS = {
    "baseline/random-walk": {
        "kind": "naive",
        "output": OUTPUT_QUANTILE_LINE,
        "display": "Random walk",
    },
    "baseline/drift": {
        "kind": "drift",
        "output": OUTPUT_QUANTILE_LINE,
        "display": "Drift",
    },
    # A bootstrap over historical candles, so there is at least one path-
    # emitting reference to hold Kronos's ghost paths against.
    "baseline/bootstrap": {
        "kind": "bootstrap",
        "output": OUTPUT_OHLCV_PATHS,
        "display": "Block bootstrap",
    },
}

# How much history the return distribution is estimated from.
LOOKBACK = 256


class BaselineAdapter(ForecastAdapter):
    """Closed-form and bootstrap baselines. No weights, no network."""

    family = "baseline"
    adapter_version = "1"

    def __init__(self, model_id: str, revision: str | None = None,
                 device: str | None = None, hardware: str | None = None):
        super().__init__(model_id, revision=revision, device=device)
        spec = BASELINE_MODELS.get(model_id)
        if spec is None:
            raise AdapterError(
                f"unknown baseline {model_id!r}; known: {sorted(BASELINE_MODELS)}")
        self.kind = spec["kind"]
        self._output = spec["output"]

    def capabilities(self) -> Capabilities:
        return Capabilities(
            output=self._output,
            hardware="cpu",
            max_context=LOOKBACK,
            asset_generality="general",
            seedable_natively=True,
        )

    def component_versions(self) -> dict[str, str]:
        # No revision to pin: the model *is* this file.
        return {"model": self.model_id, "kind": self.kind, "lookback": str(LOOKBACK)}

    def resolve_revision(self) -> str:
        self._resolved_revision = "builtin"
        return self._resolved_revision

    def load(self, model_id: str | None = None, revision: str | None = None):
        if model_id and model_id != self.model_id:
            spec = BASELINE_MODELS.get(model_id)
            if spec is None:
                raise AdapterError(f"unknown baseline {model_id!r}")
            self.model_id = model_id
            self.kind = spec["kind"]
            self._output = spec["output"]
        self.resolve_revision()
        self._model = self.kind
        return self

    def predict(self, context_ohlcv: pd.DataFrame, horizon: int,
                n_samples: int = config.DEFAULT_N_SAMPLES, seed: int = 0,
                issued_ts: pd.Timestamp | None = None) -> ForecastResult:
        check_context(context_ohlcv, issued_ts=issued_ts)
        if horizon < 1:
            raise AdapterError("horizon must be at least 1")
        if self._model is None:
            self.load()

        ctx = self._trim(context_ohlcv)
        close = ctx["close"].to_numpy(dtype="float64")
        last = float(close[-1])

        log_returns = np.diff(np.log(close))
        if len(log_returns) < 8:
            raise AdapterError("need at least 9 bars to estimate a return distribution")

        steps = np.arange(1, horizon + 1, dtype="float64")

        # A local RNG, not the global one: a baseline must not be able to
        # perturb the sampling of whatever model runs after it in a matchup.
        rng = np.random.default_rng(int(seed))

        if self.kind == "bootstrap":
            paths = self._bootstrap_paths(ctx, horizon, n_samples, rng)
            close_paths = paths[:, :, OHLCV_COLUMNS.index("close")]
            quantiles = self._quantiles_from_paths(close_paths)
            return ForecastResult(
                quantiles=quantiles, levels=config.QUANTILE_LEVELS,
                horizon=horizon, context_len=len(ctx),
                inference_version=self.inference_version(),
                seed=int(seed), n_samples=int(n_samples), paths=paths,
            )

        # Analytic band. The median is flat for `naive` and extends the mean
        # historical return for `drift`; the width is the empirical return
        # quantile scaled by sqrt(h), which is the random walk's own spread.
        mu = float(log_returns.mean()) if self.kind == "drift" else 0.0
        sigma = float(log_returns.std(ddof=1))

        cols = []
        for level in config.QUANTILE_LEVELS:
            # Normal quantile via the error function's inverse, so scipy is not
            # a dependency for the one family that has no dependencies.
            z = _norm_ppf(level)
            cols.append(last * np.exp(mu * steps + z * sigma * np.sqrt(steps)))
        quantiles = np.stack(cols, axis=1)

        return ForecastResult(
            quantiles=quantiles, levels=config.QUANTILE_LEVELS,
            horizon=horizon, context_len=len(ctx),
            inference_version=self.inference_version(),
            seed=int(seed), n_samples=int(n_samples), paths=None,
        )

    def _bootstrap_paths(self, ctx: pd.DataFrame, horizon: int,
                         n_samples: int, rng) -> np.ndarray:
        """Resample historical candles in blocks, chained off the last close.

        Blocks rather than single bars, because independent draws destroy the
        volatility clustering that makes a market path look like a market path.
        """
        frame = ctx[list(OHLCV_COLUMNS)].to_numpy(dtype="float64")
        close = frame[:, OHLCV_COLUMNS.index("close")]
        log_returns = np.diff(np.log(close))
        # Per-bar shape: how open/high/low/volume sat relative to that bar's
        # close. Resampling the shape keeps the sampled candles plausible.
        ratios = frame[1:, :4] / close[1:, None]
        volumes = frame[1:, OHLCV_COLUMNS.index("volume")]

        block = max(1, min(8, len(log_returns) // 4))
        n = len(log_returns)
        paths = np.empty((n_samples, horizon, len(OHLCV_COLUMNS)), dtype="float64")

        for s in range(n_samples):
            picked_r, picked_i = [], []
            while len(picked_r) < horizon:
                start = int(rng.integers(0, max(1, n - block)))
                for j in range(start, min(start + block, n)):
                    picked_r.append(log_returns[j])
                    picked_i.append(j)
            picked_r = np.asarray(picked_r[:horizon])
            picked_i = np.asarray(picked_i[:horizon])

            closes = float(close[-1]) * np.exp(np.cumsum(picked_r))
            shape = ratios[picked_i]
            paths[s, :, 0] = closes * shape[:, 0]     # open
            paths[s, :, 1] = closes * shape[:, 1]     # high
            paths[s, :, 2] = closes * shape[:, 2]     # low
            paths[s, :, 3] = closes                   # close
            paths[s, :, 4] = volumes[picked_i]

        # Keep the sampled candle self-consistent, same rule as Kronos.
        body_hi = np.maximum(paths[:, :, 0], paths[:, :, 3])
        body_lo = np.minimum(paths[:, :, 0], paths[:, :, 3])
        paths[:, :, 1] = np.maximum(paths[:, :, 1], body_hi)
        paths[:, :, 2] = np.minimum(paths[:, :, 2], body_lo)
        return paths


def _norm_ppf(p: float) -> float:
    """Inverse standard normal CDF, via the inverse error function."""
    from math import erf, sqrt

    if not 0.0 < p < 1.0:
        raise AdapterError(f"quantile level {p} out of range")
    # Bisection: exact enough (1e-12) and keeps this file dependency-free.
    lo, hi = -12.0, 12.0
    for _ in range(200):
        mid = (lo + hi) / 2.0
        if 0.5 * (1.0 + erf(mid / sqrt(2.0))) < p:
            lo = mid
        else:
            hi = mid
    return (lo + hi) / 2.0