backend/forecasting/providers/timesfm_provider.py

from __future__ import annotations

import asyncio
import importlib
import logging
import os
import time
from typing import Any, Dict, Optional, Tuple

import numpy as np
import pandas as pd
from fastapi import HTTPException

try:
    import torch
    TORCH_IMPORT_ERROR: Optional[str] = None
except Exception as exc:  # pragma: no cover - environment dependent
    torch = None  # type: ignore[assignment]
    TORCH_IMPORT_ERROR = str(exc)


TIMESFM_AVAILABLE = False
TIMESFM_IMPORT_ERROR: Optional[str] = None
timesfm = None

if TORCH_IMPORT_ERROR:
    TIMESFM_IMPORT_ERROR = TORCH_IMPORT_ERROR
else:
    try:
        import timesfm as _timesfm  # type: ignore[import-not-found]

        timesfm = _timesfm
        TIMESFM_AVAILABLE = True
    except Exception as exc:  # pragma: no cover - import environment dependent
        TIMESFM_IMPORT_ERROR = str(exc)


_TIMESFM_RUNTIME_PATCHED = False


def _patch_timesfm_runtime_compatibility(logger: logging.Logger) -> None:
    global _TIMESFM_RUNTIME_PATCHED
    if _TIMESFM_RUNTIME_PATCHED or timesfm is None:
        return

    try:
        internal_module = importlib.import_module("timesfm.timesfm_2p5.timesfm_2p5_torch")
        model_module_cls = getattr(internal_module, "TimesFM_2p5_200M_torch_module", None)
        if model_module_cls is None or getattr(model_module_cls, "_aiforecast_meta_patch", False):
            _TIMESFM_RUNTIME_PATCHED = True
            return

        def _patched_load_checkpoint(self: Any, path: str, **kwargs: Any) -> None:
            tensors = internal_module.load_file(path)
            has_meta_parameters = any(
                getattr(parameter, "is_meta", False)
                for parameter in self.parameters()
            )
            try:
                if has_meta_parameters:
                    self.load_state_dict(tensors, strict=True, assign=True)
                else:
                    self.load_state_dict(tensors, strict=True)
            except TypeError:
                if has_meta_parameters:
                    self.to_empty(device=self.device)
                self.load_state_dict(tensors, strict=True)

            self.to(self.device)
            torch_compile = kwargs.get("torch_compile", True)
            if torch_compile:
                internal_module.logging.info("Compiling model...")
                self = torch.compile(self)
            self.eval()

        model_module_cls.load_checkpoint = _patched_load_checkpoint
        model_module_cls._aiforecast_meta_patch = True
        _TIMESFM_RUNTIME_PATCHED = True
    except Exception as exc:  # pragma: no cover - defensive guard
        logger.warning("[TimesFM] Runtime compatibility patch skipped: %s", exc)


class TimesFMForecaster:
    key = "timesfm"
    label = "TimesFM"

    MODEL_HF_ID = os.getenv("TIMESFM_MODEL_HF_ID", "google/timesfm-2.5-200m-pytorch")
    MAX_CONTEXT = 8_192
    MAX_HORIZON = 300
    _Q10 = 1
    _Q50 = 5
    _Q90 = 9

    def __init__(self, logger: logging.Logger | None = None) -> None:
        self._logger = logger or logging.getLogger("ai-forecast.timesfm")
        self._model: Optional[Any] = None
        self._loaded = False
        self._lock: Optional[asyncio.Lock] = None
        self._predict_lock: Optional[asyncio.Lock] = None
        self._compiled_horizon: Optional[int] = None

    @property
    def available(self) -> bool:
        return TIMESFM_AVAILABLE

    @property
    def import_error(self) -> Optional[str]:
        return TIMESFM_IMPORT_ERROR

    @staticmethod
    def _prepare_feature_frame(df: pd.DataFrame) -> pd.DataFrame:
        ohlc4 = (
            df[["open", "high", "low", "close"]]
            .mean(axis=1)
            .astype(np.float32)
        )
        return pd.DataFrame({"ohlc4": ohlc4})

    @classmethod
    def _extract_ohlc4_series(cls, df: pd.DataFrame) -> np.ndarray:
        feature_frame = cls._prepare_feature_frame(df)
        ohlc4_series = feature_frame["ohlc4"].to_numpy(dtype=np.float32, copy=False)
        cls._validate_input_series(ohlc4_series)
        return ohlc4_series

    @staticmethod
    def _validate_input_series(series: np.ndarray) -> None:
        if series.ndim != 1:
            raise HTTPException(status_code=422, detail="TimesFM input series must be 1-D")
        if len(series) < 32:
            raise HTTPException(status_code=422, detail="TimesFM requires at least 32 OHLC4 points")
        if not np.isfinite(series).all():
            raise HTTPException(status_code=422, detail="TimesFM input contains non-finite OHLC4 values")

    @classmethod
    def _validate_output_tensors(
        cls,
        point_forecast: np.ndarray,
        quantile_forecast: np.ndarray,
        horizon: int,
    ) -> Tuple[np.ndarray, np.ndarray, np.ndarray, Dict[str, Any]]:
        point_arr = np.asarray(point_forecast)
        quant_arr = np.asarray(quantile_forecast)

        if point_arr.ndim != 2 or point_arr.shape[0] != 1 or point_arr.shape[1] < horizon:
            raise HTTPException(
                status_code=500,
                detail=f"TimesFM point forecast has invalid shape: {list(point_arr.shape)}",
            )
        if (
            quant_arr.ndim != 3
            or quant_arr.shape[0] != 1
            or quant_arr.shape[1] < horizon
            or quant_arr.shape[2] <= cls._Q90
        ):
            raise HTTPException(
                status_code=500,
                detail=f"TimesFM quantile forecast has invalid shape: {list(quant_arr.shape)}",
            )

        p10 = quant_arr[0, :horizon, cls._Q10].astype(float)
        p50 = point_arr[0, :horizon].astype(float)
        p50_from_quantiles = quant_arr[0, :horizon, cls._Q50].astype(float)
        p90 = quant_arr[0, :horizon, cls._Q90].astype(float)

        if not (np.isfinite(p10).all() and np.isfinite(p50).all() and np.isfinite(p90).all()):
            raise HTTPException(status_code=500, detail="TimesFM output contains non-finite values")

        quantiles_monotonic = bool(
            np.all(p10 <= (p50 + 1e-6))
            and np.all(p50 <= (p90 + 1e-6))
        )
        if not quantiles_monotonic:
            raise HTTPException(status_code=500, detail="TimesFM returned non-monotonic quantiles")

        median_matches_point_forecast = bool(
            np.allclose(p50, p50_from_quantiles, atol=1e-4, rtol=1e-4)
        )
        if not median_matches_point_forecast:
            logging.getLogger("ai-forecast").warning(
                "[TimesFM] point_forecast differs from q50 quantile output"
            )

        output_validation = {
            "point_shape": [int(dim) for dim in point_arr.shape],
            "quantile_shape": [int(dim) for dim in quant_arr.shape],
            "median_matches_point_forecast": median_matches_point_forecast,
            "quantiles_monotonic": quantiles_monotonic,
        }
        return p10, p50, p90, output_validation

    async def _get_lock(self) -> asyncio.Lock:
        if self._lock is None:
            self._lock = asyncio.Lock()
        return self._lock

    async def _get_predict_lock(self) -> asyncio.Lock:
        if self._predict_lock is None:
            self._predict_lock = asyncio.Lock()
        return self._predict_lock

    @property
    def is_ready(self) -> bool:
        return self._loaded

    @property
    def device(self) -> str:
        if self._model is None:
            return "not_loaded"
        try:
            return str(self._model.model.device)
        except Exception:
            return "cpu"

    def _compile(self, horizon: int) -> None:
        output_patch = 128
        max_horizon = int(np.ceil(horizon / output_patch) * output_patch)
        max_horizon = max(max_horizon, output_patch)
        max_horizon = min(max_horizon, self.MAX_HORIZON)
        context_len = self.MAX_CONTEXT
        assert timesfm is not None
        self._model.compile(
            timesfm.ForecastConfig(
                max_context=context_len,
                max_horizon=max_horizon,
                normalize_inputs=True,
                use_continuous_quantile_head=True,
                force_flip_invariance=True,
                infer_is_positive=True,
                fix_quantile_crossing=True,
            )
        )
        compiled_config = getattr(self._model, "forecast_config", None)
        actual_ctx = int(getattr(compiled_config, "max_context", context_len))
        actual_max_horizon = int(getattr(compiled_config, "max_horizon", max_horizon))
        self._compiled_horizon = actual_max_horizon
        self._logger.info("[TimesFM] Compiled: ctx=%d max_horizon=%d", actual_ctx, actual_max_horizon)

    async def _lazy_load(self) -> None:
        if self._loaded:
            return
        lock = await self._get_lock()
        async with lock:
            if self._loaded:
                return
            if not TIMESFM_AVAILABLE or timesfm is None:
                raise HTTPException(
                    status_code=503,
                    detail="TimesFM not installed. Run: pip install timesfm[torch]",
                )
            try:
                _patch_timesfm_runtime_compatibility(self._logger)
                self._logger.info("[TimesFM] Loading %s ...", self.MODEL_HF_ID)
                model = await asyncio.to_thread(
                    timesfm.TimesFM_2p5_200M_torch.from_pretrained,
                    self.MODEL_HF_ID,
                    torch_compile=False,
                )
                self._model = model
                self._compile(self.MAX_HORIZON)
                self._loaded = True
                self._logger.info("[TimesFM] Ready on %s", self.device)
            except Exception as exc:
                self._logger.error("[TimesFM] Init failed: %s", exc, exc_info=True)
                raise HTTPException(status_code=500, detail=f"TimesFM init failed: {exc}")

    async def forecast(
        self,
        df: pd.DataFrame,
        horizon: int,
        interval: str = "",
        step_seconds: int = 0,
        symbol: str = "",
    ) -> Dict[str, Any]:
        del interval, step_seconds, symbol
        await self._lazy_load()
        if self._model is None:
            raise HTTPException(status_code=500, detail="TimesFM model is not initialized")

        try:
            ohlc4_series = self._extract_ohlc4_series(df)
            context_len = min(len(ohlc4_series), self.MAX_CONTEXT)
            ohlc4_context = ohlc4_series[-context_len:]

            if self._compiled_horizon is None or horizon > self._compiled_horizon:
                self._compile(horizon)

            started_at = time.time()
            predict_lock = await self._get_predict_lock()
            async with predict_lock:
                point_forecast, quantile_forecast = await asyncio.to_thread(
                    self._model.forecast,
                    horizon,
                    [ohlc4_context],
                )
            elapsed = time.time() - started_at
            self._logger.info(
                "[TimesFM] %.2fs | horizon=%d ctx=%d device=%s",
                elapsed,
                horizon,
                context_len,
                self.device,
            )

            if torch is not None and torch.cuda.is_available():
                torch.cuda.empty_cache()

            p10, p50, p90, output_validation = self._validate_output_tensors(
                point_forecast=point_forecast,
                quantile_forecast=quantile_forecast,
                horizon=horizon,
            )

            return {
                "p10": p10,
                "p50": p50,
                "p90": p90,
                "model_name": self.MODEL_HF_ID,
                "context_length": context_len,
                "output_horizon": horizon,
                "input_semantics": {
                    "feature_channels": ["ohlc4"],
                    "active_forecast_channels": ["ohlc4"],
                    "ignored_channels": ["volume"],
                    "price_mode": "ohlc4_single_channel",
                    "base_signal": "ohlc4",
                    "volume_mode": "omitted",
                    "amount_mode": "omitted",
                    "adapter_mode": "timesfm_native",
                    "normalization": "timesfm_internal_revin",
                },
                "input_validation": {
                    "series_field": "ohlc4",
                    "dtype": str(ohlc4_context.dtype),
                    "is_1d": True,
                    "finite": True,
                },
                "output_semantics": {
                    "forecast_channel": "ohlc4",
                    "forecast_mode": "single_future_ohlc4_line",
                    "quantile_fields": ["p10", "p50", "p90"],
                    "candle_projection": "omitted",
                    "reference_baseline": "last_ohlc4",
                },
                "output_validation": output_validation,
            }
        except HTTPException:
            raise
        except Exception as exc:
            self._logger.error("[TimesFM] Forecast failed: %s", exc, exc_info=True)
            raise HTTPException(status_code=500, detail=f"TimesFM prediction failed: {exc}")