src/adapters/base.py

"""The adapter interface every forecasting model is reached through.

Three ideas carry most of the weight here.

**Capabilities are declared, not inferred.** An adapter says what kind of
output it produces, what hardware it needs, and how much context it can take.
The UI renders strictly from that declaration, so there is no place outside
this package where a model is special-cased by name. Adding a family is a new
module plus a registry entry, never an `if model == ...` in a renderer.

**Seeding is part of the interface.** `predict` takes a seed and is required to
be bit-reproducible under it. Some upstream models expose no seed argument at
all; those are wrapped so that their sampling is seeded anyway. A forecast that
cannot be reproduced cannot be audited, and an archive of unauditable forecasts
is not a track record.

**Contexts are validated, not trusted.** `check_context` runs on every call. It
is the structural guarantee behind the no-lookahead rule: a context bar dated
after the issue timestamp raises rather than quietly producing a forecast that
would look brilliant.
"""

from __future__ import annotations

import hashlib
import logging
from abc import ABC, abstractmethod
from dataclasses import dataclass, field

import numpy as np
import pandas as pd

from .. import config

log = logging.getLogger("arena.adapters")

OHLCV_COLUMNS = ("open", "high", "low", "close", "volume")

# Output kinds. `ohlcv_paths` models emit whole sampled candlestick paths and
# can drive the ghost-path and dispersion views; `quantile_line` models emit a
# band on the close and nothing else.
OUTPUT_QUANTILE_LINE = "quantile_line"
OUTPUT_OHLCV_PATHS = "ohlcv_paths"
OUTPUT_KINDS = (OUTPUT_QUANTILE_LINE, OUTPUT_OHLCV_PATHS)

HARDWARE_CPU = "cpu"
HARDWARE_GPU = "gpu"
HARDWARE_TIERS = (HARDWARE_CPU, HARDWARE_GPU)


class AdapterError(RuntimeError):
    """Anything an adapter refuses to do."""


class ModelNotAllowed(AdapterError):
    """The requested adapter family is not on the allow-list."""


class LookaheadError(AdapterError):
    """A context window reached past the moment the forecast is issued."""


class ContextError(AdapterError):
    """The context window is unusable -- gaps, NaNs, or too short."""


# --------------------------------------------------------------------------
# Capabilities
# --------------------------------------------------------------------------


@dataclass(frozen=True)
class Capabilities:
    """What a model can do, as the model itself declares it.

    `asset_generality` is a plain-language claim about training scope, not a
    quality score: "financial" means the model was pre-trained on market data,
    "general" means arbitrary time series. It drives one honest caveat in the
    UI and nothing else.
    """

    output: str
    hardware: str
    max_context: int
    asset_generality: str = "general"
    seedable_natively: bool = True

    def __post_init__(self):
        if self.output not in OUTPUT_KINDS:
            raise AdapterError(f"unknown output kind {self.output!r}")
        if self.hardware not in HARDWARE_TIERS:
            raise AdapterError(f"unknown hardware tier {self.hardware!r}")
        if self.max_context < 1:
            raise AdapterError("max_context must be positive")

    @property
    def emits_paths(self) -> bool:
        return self.output == OUTPUT_OHLCV_PATHS

    def as_dict(self) -> dict:
        return {
            "output": self.output,
            "hardware": self.hardware,
            "max_context": self.max_context,
            "asset_generality": self.asset_generality,
            "seedable_natively": self.seedable_natively,
        }


# --------------------------------------------------------------------------
# Result
# --------------------------------------------------------------------------


@dataclass
class ForecastResult:
    """One multi-step forecast.

    `quantiles` is (horizon, n_levels) over the *close* price, always present.
    `paths` is (n_samples, horizon, 5) of OHLCV and is present only for
    `ohlcv_paths` models -- the UI checks the capability, never this field, so
    that a model which merely happens to return paths cannot change how it is
    rendered.
    """

    quantiles: np.ndarray
    levels: tuple[float, ...]
    horizon: int
    context_len: int
    inference_version: str
    seed: int
    n_samples: int
    paths: np.ndarray | None = None

    def __post_init__(self):
        self.quantiles = np.asarray(self.quantiles, dtype="float64")
        if self.quantiles.shape != (self.horizon, len(self.levels)):
            raise AdapterError(
                f"quantiles shape {self.quantiles.shape} != "
                f"{(self.horizon, len(self.levels))}"
            )
        if not np.all(np.isfinite(self.quantiles)):
            raise AdapterError("forecast contains non-finite values")
        # Quantiles must not cross. Sorting along the level axis is the honest
        # repair for the tiny numerical inversions that sampling produces, and
        # it is what makes the coverage arithmetic downstream well-defined.
        self.quantiles = np.sort(self.quantiles, axis=1)
        if self.paths is not None:
            self.paths = np.asarray(self.paths, dtype="float64")
            if self.paths.ndim != 3 or self.paths.shape[1] != self.horizon:
                raise AdapterError(f"paths shape {self.paths.shape} is not (n, {self.horizon}, k)")

    def level_index(self, level: float) -> int:
        for i, lv in enumerate(self.levels):
            if abs(lv - level) < 1e-9:
                return i
        raise AdapterError(f"level {level} not in {self.levels}")

    def band(self, low: float = 0.1, high: float = 0.9):
        return (self.quantiles[:, self.level_index(low)],
                self.quantiles[:, self.level_index(high)])

    def median(self) -> np.ndarray:
        return self.quantiles[:, self.level_index(0.5)]

    def dispersion(self) -> np.ndarray:
        """Per-step spread, as a fraction of the median.

        Sampled-path standard deviation where paths exist, band half-width
        otherwise. Both are a width; they are not the same statistic, and the
        UI labels which one it is showing.
        """
        med = np.abs(self.median()) + 1e-12
        if self.paths is not None:
            close = self.paths[:, :, OHLCV_COLUMNS.index("close")]
            return close.std(axis=0) / med
        lo, hi = self.band()
        return (hi - lo) / 2.0 / med


# --------------------------------------------------------------------------
# Context validation
# --------------------------------------------------------------------------


def check_context(context: pd.DataFrame, issued_ts: pd.Timestamp | None = None,
                  min_len: int = 32) -> pd.DataFrame:
    """Validate a context window, or raise.

    This is the structural half of the no-lookahead guarantee. It is not a
    convention that callers are asked to honour: every adapter runs it on every
    call, so a forecast issued from data it should not have seen fails loudly
    at the point of use.
    """
    if not isinstance(context, pd.DataFrame):
        raise ContextError("context must be a DataFrame")

    missing = [c for c in OHLCV_COLUMNS if c not in context.columns]
    if missing:
        raise ContextError(f"context is missing columns {missing}")

    if len(context) < min_len:
        raise ContextError(f"context has {len(context)} bars, need at least {min_len}")

    if "ts" not in context.columns:
        raise ContextError("context must carry a 'ts' column")

    ts = pd.to_datetime(context["ts"], utc=True)
    if ts.isna().any():
        raise ContextError("context has unparseable timestamps")
    if not ts.is_monotonic_increasing:
        raise ContextError("context timestamps are not sorted ascending")
    if ts.duplicated().any():
        raise ContextError("context has duplicate timestamps")

    prices = context[list(OHLCV_COLUMNS)]
    if not np.isfinite(prices.to_numpy(dtype="float64")).all():
        raise ContextError("context contains NaN or infinite values")
    if (context[["open", "high", "low", "close"]].to_numpy(dtype="float64") <= 0).any():
        raise ContextError("context contains non-positive prices")

    # A hole in the cache means the model sees two bars as adjacent when they
    # are weeks apart, and produces a confident wrong answer from it.
    #
    # But "irregular spacing" is not the same as "missing data". An hourly
    # equity series closes for ~17 hours every night and ~65 hours every
    # weekend; those gaps are the market, not the cache. Judging them against
    # the modal bar rejected every SPY and NVDA hourly context outright.
    #
    # So the threshold adapts to the series' own gap distribution: a session
    # boundary recurs and therefore sits inside the 99th percentile, while a
    # genuinely absent period stands outside it.
    if len(ts) > 20:
        deltas = ts.diff().dropna()
        modal = deltas.mode()
        if len(modal):
            step = modal.iloc[0]
            p99 = deltas.quantile(0.99)

            # Any *single* gap far beyond the series' own worst regular one is
            # a hole, however few there are. This is the check that survives a
            # series where absence has become the pattern -- an adaptive
            # threshold alone would quietly normalise that.
            hard = max(step * 24, p99 * 5)
            worst = deltas.max()
            if worst > hard:
                raise ContextError(
                    f"context has a {worst} gap, far beyond its own session "
                    f"pattern; the cache is incomplete for this window"
                )

            # And an accumulation of smaller anomalies is a hole too.
            ceiling = max(step * 4, p99 * 1.5)
            bad = deltas[deltas > ceiling]
            if len(bad) > max(2, len(ts) // 50):
                raise ContextError(
                    f"context has {len(bad)} gaps beyond its own session pattern "
                    f"(> {ceiling}); the cache is incomplete for this window"
                )

    if issued_ts is not None:
        issued = pd.Timestamp(issued_ts)
        if issued.tzinfo is None:
            issued = issued.tz_localize("UTC")
        last = ts.iloc[-1]
        if last > issued:
            raise LookaheadError(
                f"context ends at {last.isoformat()}, after issued_ts "
                f"{issued.isoformat()}: a forecast may only see data at or "
                f"before the moment it is issued"
            )

    return context


# --------------------------------------------------------------------------
# Determinism
# --------------------------------------------------------------------------


def seed_everything(seed: int) -> None:
    """Pin every RNG an adapter might reach for.

    Seeding numpy alone is not enough: the sampling models draw through torch,
    and Kronos in particular calls `torch.multinomial` with no seed argument of
    its own. This is what makes "same inputs and seed produce bit-identical
    output" true rather than aspirational.
    """
    seed = int(seed) % (2 ** 31 - 1)
    np.random.seed(seed)
    try:
        import random

        random.seed(seed)
    except Exception:  # pragma: no cover
        pass
    try:
        import torch

        torch.manual_seed(seed)
        if torch.cuda.is_available():
            torch.cuda.manual_seed_all(seed)
    except ImportError:
        pass


# --------------------------------------------------------------------------
# The interface
# --------------------------------------------------------------------------


class ForecastAdapter(ABC):
    """Uniform interface over multi-step probabilistic forecasters."""

    family: str = "base"
    # Bumped when this adapter's own logic changes the numbers it produces.
    adapter_version: str = "1"

    def __init__(self, model_id: str, revision: str | None = None,
                 device: str | None = None):
        self.model_id = model_id
        self.revision = revision
        self._device = device
        self._model = None
        self._resolved_revision: str | None = None

    @property
    def device(self) -> str:
        """Resolved on first use, never at construction.

        On ZeroGPU, `torch.cuda.is_available()` raises unless it is called
        inside a `@spaces.GPU` function -- so probing the device eagerly made
        merely *constructing* an adapter fatal, including for CPU-tier models
        that never wanted a GPU. Deferring it means construction is free and
        the probe happens inside the GPU call, where it is legal.
        """
        if self._device is None:
            self._device = default_device()
        return self._device

    # -- interface --------------------------------------------------------

    @abstractmethod
    def load(self, model_id: str | None = None, revision: str | None = None) -> "ForecastAdapter":
        """Materialise the model. Idempotent."""

    @abstractmethod
    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:
        """Forecast `horizon` bars ahead from `context_ohlcv`.

        Must be bit-reproducible in `seed`.
        """

    @abstractmethod
    def capabilities(self) -> Capabilities:
        """What this model can do. The UI renders from this and nothing else."""

    # -- shared -----------------------------------------------------------

    def component_versions(self) -> dict[str, str]:
        """Every external thing whose version changes the output.

        Overridden by families that load a companion repo (Kronos ships its
        tokenizer separately) or depend on an inference package whose version
        moves the numbers.
        """
        return {"model": f"{self.model_id}@{self.resolved_revision}"}

    def inference_version(self) -> str:
        """Identity of everything that determines the output values.

        A short hash rather than the full component list, because it is written
        onto every archived row; the components themselves are recorded once in
        the registry so the hash can always be expanded.
        """
        parts = [config.INFERENCE_VERSION, self.family, self.adapter_version]
        parts += [f"{k}={v}" for k, v in sorted(self.component_versions().items())]
        digest = hashlib.sha256("|".join(parts).encode()).hexdigest()[:12]
        return f"{config.INFERENCE_VERSION}+{self.family}.{digest}"

    @property
    def resolved_revision(self) -> str:
        return self._resolved_revision or self.revision or "unpinned"

    def resolve_revision(self) -> str:
        """Pin the model to an immutable commit sha before any inference runs.

        A floating `main` would mean two forecasts a week apart are not
        comparable and neither is reproducible, which defeats the archive.
        """
        if self._resolved_revision:
            return self._resolved_revision
        from huggingface_hub import HfApi

        info = HfApi().model_info(self.model_id, revision=self.revision)
        self._resolved_revision = info.sha
        return self._resolved_revision

    # -- helpers for subclasses -------------------------------------------

    def _trim(self, context: pd.DataFrame) -> pd.DataFrame:
        """Cut a context down to what the model can actually attend over."""
        cap = self.capabilities().max_context
        return context.iloc[-cap:] if len(context) > cap else context

    @staticmethod
    def _quantiles_from_paths(paths_close: np.ndarray,
                              levels=config.QUANTILE_LEVELS) -> np.ndarray:
        """(n_samples, horizon) -> (horizon, n_levels)."""
        return np.quantile(paths_close, list(levels), axis=0).T


def default_device() -> str:
    """The device to run on, when the caller has not said.

    On ZeroGPU this never probes. `torch.cuda.is_available()` triggers a
    low-level CUDA init that ZeroGPU forbids outside a `@spaces.GPU` function,
    and it does not fail politely -- it takes the load down. Inside a GPU
    function the device is known to be cuda anyway, so `runtime` passes it
    explicitly and this is only the fallback for everywhere else.
    """
    try:
        import spaces  # noqa: F401 - presence is the signal

        return "cpu"
    except ImportError:
        pass
    try:
        import torch

        if torch.cuda.is_available():
            return "cuda"
    except Exception:
        pass
    return "cpu"