src/adapters.py

"""Forecast model adapters.

Adding a model means giving a model id and naming an adapter *family*. The
family list is fixed in `config.ALLOWED_ADAPTER_FAMILIES`; nothing here ever
imports, downloads or executes code chosen by a user. A user-supplied model id
is loaded only through an already-vetted family's loader, and `trust_remote_code`
is never enabled.

Every adapter reports an `inference_version` that pins the adapter's own logic
alongside the model revision, so a stored slice can always be traced back to
exactly what produced it.
"""

from __future__ import annotations

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

import numpy as np
import pandas as pd

from . import config

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

DEFAULT_QUANTILES = (0.1, 0.5, 0.9)


class AdapterError(RuntimeError):
    pass


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


@dataclass
class Forecast:
    """One horizon-1 forecast per input window."""

    q10: np.ndarray
    q50: np.ndarray
    q90: np.ndarray
    context_len: int

    def __post_init__(self):
        if not (len(self.q10) == len(self.q50) == len(self.q90)):
            raise AdapterError("quantile arrays have mismatched lengths")

    def as_frame(self, ts: pd.DatetimeIndex, inference_version: str) -> pd.DataFrame:
        if len(ts) != len(self.q50):
            raise AdapterError(
                f"timestamp count {len(ts)} != forecast count {len(self.q50)}"
            )
        # Quantiles must not cross; sorting is the honest repair for tiny
        # numerical inversions and makes the store's validator pass.
        stacked = np.sort(np.vstack([self.q10, self.q50, self.q90]), axis=0)
        return pd.DataFrame({
            "ts": ts,
            "q10": stacked[0], "q50": stacked[1], "q90": stacked[2],
            "context_len": self.context_len,
            "inference_version": inference_version,
        })


class ForecastAdapter(ABC):
    """Uniform interface over quantile forecasters."""

    family: str = "base"
    adapter_version: str = "1"

    def __init__(self, model_id: str, revision: str | None = None,
                 context_len: int = 512, device: str | None = None):
        self.model_id = model_id
        self.revision = revision
        self.context_len = context_len
        self.device = device or _default_device()
        self._model = None
        self._resolved_revision: str | None = None

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

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

    @abstractmethod
    def predict(self, context_windows: np.ndarray) -> Forecast:
        """`context_windows` is (n_windows, context_len). Returns horizon-1 quantiles."""

    def inference_version(self) -> str:
        """Identity of everything that determines the output values."""
        rev = self._resolved_revision or self.revision or "unpinned"
        payload = f"{self.family}|{self.adapter_version}|{self.model_id}|{rev}|{self.context_len}"
        digest = hashlib.sha256(payload.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"

    # -- shared helpers ---------------------------------------------------

    def resolve_revision(self) -> str:
        """Pin the model to an immutable commit sha before any inference runs."""
        if self._resolved_revision:
            return self._resolved_revision
        try:
            from huggingface_hub import HfApi

            info = HfApi().model_info(self.model_id, revision=self.revision)
            self._resolved_revision = info.sha
        except Exception as e:
            log.warning("could not resolve revision for %s: %s", self.model_id, e)
            self._resolved_revision = self.revision or "unpinned"
        return self._resolved_revision


def _default_device() -> str:
    try:
        import torch

        if torch.cuda.is_available():
            return "cuda"
        if getattr(torch.backends, "mps", None) and torch.backends.mps.is_available():
            return "mps"
    except Exception:
        pass
    return "cpu"


def _release(device: str) -> None:
    """Free accelerator memory between inference chunks."""
    try:
        import torch

        if device == "cuda":
            torch.cuda.empty_cache()
        elif device == "mps" and hasattr(torch, "mps"):
            torch.mps.empty_cache()
    except Exception:
        pass


# --------------------------------------------------------------------------
# Chronos / Chronos-Bolt
# --------------------------------------------------------------------------


class ChronosAdapter(ForecastAdapter):
    """amazon/chronos-* and amazon/chronos-bolt-* quantile forecasters."""

    family = "chronos"
    adapter_version = "1"

    def load(self, model_id: str | None = None, revision: str | None = None):
        if model_id:
            self.model_id = model_id
        if revision:
            self.revision = revision
        if self._model is not None:
            return self

        self.resolve_revision()
        try:
            from chronos import BaseChronosPipeline
        except ImportError as e:
            raise AdapterError(
                "chronos-forecasting is not installed; add it to requirements.txt"
            ) from e

        import torch

        dtype = torch.float32 if self.device in ("cpu", "mps") else torch.bfloat16
        self._model = BaseChronosPipeline.from_pretrained(
            self.model_id,
            revision=self._resolved_revision if self._resolved_revision != "unpinned" else None,
            device_map=self.device,
            torch_dtype=dtype,
        )
        return self

    # Bolt emits quantiles directly and is happy with wide batches. The
    # original T5 Chronos *samples* forecast paths instead, so its peak memory
    # is batch x num_samples x context and a wide batch OOMs a GPU outright.
    # Chunking here rather than at the call site means every caller -- seed
    # script, ZeroGPU function, tests -- is protected without knowing which
    # variant it holds.
    BOLT_CHUNK = 256
    T5_CHUNK = 16
    T5_NUM_SAMPLES = 20
    CHRONOS2_CHUNK = 128

    @property
    def _is_bolt(self) -> bool:
        return "bolt" in (self.model_id or "").lower()

    @property
    def _is_chronos2(self) -> bool:
        """Chronos-2 (`amazon/chronos-2`).

        Detected from the id rather than from the loaded pipeline class so the
        chunk size is known before the weights are fetched.
        """
        name = (self.model_id or "").lower()
        return "chronos-2" in name or "chronos2" in name

    @property
    def chunk_size(self) -> int:
        if self._is_chronos2:
            return self.CHRONOS2_CHUNK
        return self.BOLT_CHUNK if self._is_bolt else self.T5_CHUNK

    def predict(self, context_windows: np.ndarray) -> Forecast:
        if self._model is None:
            self.load()
        import torch

        ctx = np.asarray(context_windows, dtype="float32")
        if ctx.ndim == 1:
            ctx = ctx[None, :]

        q_levels = list(DEFAULT_QUANTILES)
        # Only the original T5 Chronos samples paths; Bolt and Chronos-2 emit
        # quantiles directly.
        if self._is_bolt or self._is_chronos2:
            extra = {}
        else:
            extra = {"num_samples": self.T5_NUM_SAMPLES}
        step = max(1, self.chunk_size)
        parts = []

        for start in range(0, len(ctx), step):
            block = ctx[start:start + step]
            tensors = [torch.tensor(row) for row in block]
            with torch.inference_mode():
                quantiles, _mean = self._model.predict_quantiles(
                    tensors, prediction_length=1, quantile_levels=q_levels, **extra,
                )
            parts.append(self._to_array(quantiles))
            del quantiles, tensors
            _release(self.device)

        arr = np.vstack(parts)
        return Forecast(q10=arr[:, 0], q50=arr[:, 1], q90=arr[:, 2],
                        context_len=ctx.shape[1])

    @staticmethod
    def _to_array(quantiles) -> np.ndarray:
        """Normalise a `predict_quantiles` result to `(batch, n_quantiles)`.

        The two shapes differ and the difference is not cosmetic:

          Bolt / T5      one stacked tensor, `(batch, horizon, quantiles)`
          Chronos-2      a *list* of per-item tensors, each
                         `(n_variates, horizon, quantiles)`

        Calling `.float()` on the list raises `AttributeError`, which is
        exactly what `amazon/chronos-2` did before this existed. Horizon is
        always 1 here, and these are univariate price series, so the leading
        variate axis is taken at index 0.
        """
        if isinstance(quantiles, (list, tuple)):
            rows = []
            for item in quantiles:
                array = ChronosAdapter._as_numpy(item)
                # (n_variates, horizon, quantiles) -> (quantiles,)
                while array.ndim > 1:
                    array = array[0]
                rows.append(array)
            return np.vstack(rows)

        return ChronosAdapter._as_numpy(quantiles)[:, 0, :]

    @staticmethod
    def _as_numpy(value) -> np.ndarray:
        """Tensor -> ndarray, by duck typing rather than an isinstance check.

        Deliberately does not import torch: the test suite runs offline with no
        GPU stack installed, and a shape-normalising helper that cannot be
        tested without a 2GB dependency would not get tested.
        """
        if hasattr(value, "float") and hasattr(value, "cpu"):
            return np.asarray(value.float().cpu().numpy())
        if isinstance(value, np.ndarray):
            return value
        raise AdapterError(
            f"unexpected predict_quantiles result: {type(value).__name__}")


# --------------------------------------------------------------------------
# TimesFM
# --------------------------------------------------------------------------


class TimesFMAdapter(ForecastAdapter):
    """google/timesfm-* forecasters.

    TimesFM returns a fixed quantile grid; the 0.1/0.5/0.9 columns are selected
    from it rather than re-derived, so the stored numbers are the model's own.
    """

    family = "timesfm"
    adapter_version = "1"
    _QUANTILE_GRID = (0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9)

    def load(self, model_id: str | None = None, revision: str | None = None):
        if model_id:
            self.model_id = model_id
        if revision:
            self.revision = revision
        if self._model is not None:
            return self

        self.resolve_revision()
        try:
            import timesfm
        except ImportError as e:
            raise AdapterError(
                "timesfm is not installed; add it to requirements.txt to enable "
                "this adapter family"
            ) from e

        backend = {"cuda": "gpu", "mps": "cpu"}.get(self.device, "cpu")
        self._model = timesfm.TimesFm(
            hparams=timesfm.TimesFmHparams(
                backend=backend, per_core_batch_size=32,
                context_len=self.context_len, horizon_len=1,
            ),
            checkpoint=timesfm.TimesFmCheckpoint(huggingface_repo_id=self.model_id),
        )
        return self

    def predict(self, context_windows: np.ndarray) -> Forecast:
        if self._model is None:
            self.load()

        ctx = np.asarray(context_windows, dtype="float32")
        if ctx.ndim == 1:
            ctx = ctx[None, :]

        _point, quantile_out = self._model.forecast(
            [row for row in ctx], freq=[0] * len(ctx)
        )
        arr = np.asarray(quantile_out)[:, 0, :]  # horizon 1
        grid = list(self._QUANTILE_GRID)
        # Column 0 of TimesFM's output is the mean; quantiles follow.
        offset = arr.shape[1] - len(grid)
        i10, i50, i90 = (grid.index(q) + offset for q in DEFAULT_QUANTILES)
        return Forecast(q10=arr[:, i10], q50=arr[:, i50], q90=arr[:, i90],
                        context_len=ctx.shape[1])


# --------------------------------------------------------------------------
# Naive baselines
# --------------------------------------------------------------------------


class BaselineAdapter(ForecastAdapter):
    """Classical forecasting baselines, exposed as first-class models.

    These exist so the leaderboard always carries a floor. A large pretrained
    forecaster that cannot beat "tomorrow looks like today" has not earned its
    inference cost, and burying that comparison would be the single easiest way
    to make this whole Space misleading.

    The method is selected by the model id: `baseline/naive`, `baseline/drift`,
    or `baseline/seasonal`. All three are deterministic and causal -- each uses
    only the trailing window, which ends at the bar being forecast.

    The interval is a Gaussian band around the point forecast, scaled by the
    window's own realised volatility, so coverage is comparable with a
    learned model's q10-q90 rather than arbitrarily wide.
    """

    family = "baseline"
    adapter_version = "1"
    Z90 = 1.2815515655446004  # normal 90th percentile

    def __init__(self, model_id: str, revision: str | None = None,
                 context_len: int = 128, device: str | None = None,
                 season: int = 7):
        super().__init__(model_id, revision, context_len, device)
        self.season = season
        self.method = (model_id.split("/")[-1] or "naive").lower()
        if self.method not in ("naive", "drift", "seasonal"):
            raise AdapterError(
                f"unknown baseline method {self.method!r}; "
                "use baseline/naive, baseline/drift or baseline/seasonal"
            )

    def load(self, model_id: str | None = None, revision: str | None = None):
        if model_id:
            self.model_id = model_id
            self.method = model_id.split("/")[-1].lower()
        # Baselines have no weights, so the "revision" is the adapter's own
        # version -- still pinned, still reproducible.
        self._resolved_revision = f"baseline-{self.adapter_version}"
        return self

    def resolve_revision(self) -> str:
        self._resolved_revision = f"baseline-{self.adapter_version}"
        return self._resolved_revision

    def predict(self, context_windows: np.ndarray) -> Forecast:
        ctx = np.asarray(context_windows, dtype="float64")
        if ctx.ndim == 1:
            ctx = ctx[None, :]
        last = ctx[:, -1]

        if self.method == "naive":
            point = last
        elif self.method == "drift":
            # Average per-step change across the window, added once.
            steps = ctx.shape[1] - 1
            point = last + (ctx[:, -1] - ctx[:, 0]) / max(steps, 1)
        else:  # seasonal
            k = min(self.season, ctx.shape[1])
            point = ctx[:, -k]

        # One-step-ahead uncertainty from the window's own step volatility.
        sigma = np.std(np.diff(ctx, axis=1), axis=1)
        sigma = np.where(np.isfinite(sigma) & (sigma > 0), sigma,
                         np.abs(last) * 1e-4)
        band = self.Z90 * sigma
        return Forecast(q10=point - band, q50=point, q90=point + band,
                        context_len=ctx.shape[1])


# --------------------------------------------------------------------------
# Placeholder (no GPU / no model available)
# --------------------------------------------------------------------------


class PlaceholderAdapter(ForecastAdapter):
    """Structurally identical synthetic output, permanently labelled.

    Exists so the UI has complete shape before every cell has real coverage.
    Its `inference_version` is the literal string `PLACEHOLDER`, which the store
    and the UI both key off to mark the data as not-real. It is deterministic:
    the same window always yields the same numbers.
    """

    family = "placeholder"
    adapter_version = "1"

    def load(self, model_id: str | None = None, revision: str | None = None):
        if model_id:
            self.model_id = model_id
        self._resolved_revision = "PLACEHOLDER"
        return self

    def inference_version(self) -> str:
        return config.PLACEHOLDER_VERSION

    def predict(self, context_windows: np.ndarray) -> Forecast:
        ctx = np.asarray(context_windows, dtype="float64")
        if ctx.ndim == 1:
            ctx = ctx[None, :]

        last = ctx[:, -1]
        # Deterministic pseudo-drift seeded by the window itself, plus a band
        # scaled to that window's realised volatility.
        seedvals = np.abs(np.sum(ctx[:, -8:], axis=1) * 1e6).astype("int64")
        drift = np.array([
            (np.random.default_rng(int(s) % (2**32)).normal(0.0, 1.0)) for s in seedvals
        ])
        vol = np.std(np.diff(ctx, axis=1), axis=1) / np.maximum(np.abs(last), 1e-9)
        vol = np.clip(vol, 1e-4, 0.2)

        q50 = last * (1.0 + 0.25 * drift * vol)
        band = last * vol * 1.28
        return Forecast(q10=q50 - band, q50=q50, q90=q50 + band,
                        context_len=ctx.shape[1])


# --------------------------------------------------------------------------
# Factory
# --------------------------------------------------------------------------

_FAMILIES: dict[str, type[ForecastAdapter]] = {
    "chronos": ChronosAdapter,
    "timesfm": TimesFMAdapter,
    "baseline": BaselineAdapter,
    "placeholder": PlaceholderAdapter,
}


def get_adapter(family: str, model_id: str, revision: str | None = None,
                context_len: int = 512, device: str | None = None) -> ForecastAdapter:
    """Build an adapter. Only allow-listed families are constructible."""
    fam = (family or "").strip().lower()
    if fam == "placeholder":
        return PlaceholderAdapter(model_id, revision, context_len, device)
    if fam not in config.ALLOWED_ADAPTER_FAMILIES:
        raise ModelNotAllowed(
            f"adapter family {family!r} is not allowed; pick one of "
            f"{list(config.ALLOWED_ADAPTER_FAMILIES)}"
        )
    return _FAMILIES[fam](model_id, revision, context_len, device)


def validate_model_id(model_id: str) -> str:
    """Reject anything that is not a plain `owner/name` Hub id."""
    mid = (model_id or "").strip()
    if not mid or mid.count("/") != 1:
        raise AdapterError(f"{model_id!r} is not a valid Hub model id (owner/name)")
    owner, name = mid.split("/")
    ok = set("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789-_.")
    if not owner or not name or set(mid) - ok - {"/"}:
        raise AdapterError(f"{model_id!r} contains characters that are not allowed")
    if ".." in mid:
        raise AdapterError("model id may not contain '..'")
    return mid


# --------------------------------------------------------------------------
# Rolling-window construction (shared by the seed script and the GPU path)
# --------------------------------------------------------------------------


def build_windows(series: pd.Series, context_len: int, stride: int = 1):
    """Yield (timestamp, trailing window) pairs.

    The window for timestamp `t` ends at `t` inclusive, so the forecast stored
    at `t` used only data available at `t`. The engine then shifts it before any
    trade can act on it.
    """
    values = series.to_numpy(dtype="float64")
    index = series.index
    n = len(values)
    if n <= context_len:
        return [], np.empty((0, context_len))

    stamps, rows = [], []
    for i in range(context_len, n, stride):
        stamps.append(index[i])
        rows.append(values[i - context_len + 1: i + 1])
    if not rows:
        return [], np.empty((0, context_len))
    return pd.DatetimeIndex(stamps), np.vstack(rows)