backend/forecasting/providers/chronos_provider.py

from __future__ import annotations

import asyncio
import hashlib
import logging
import os
import random
import sys
import time
from pathlib import Path
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)


PROJECT_ROOT = Path(__file__).resolve().parents[3]
CHRONOS_VENDOR_SRC = PROJECT_ROOT / "libs" / "chronos-forecasting" / "src"

CHRONOS_AVAILABLE = False
CHRONOS_IMPORT_ERROR: Optional[str] = None
ChronosPipeline = None

if TORCH_IMPORT_ERROR:
    CHRONOS_IMPORT_ERROR = TORCH_IMPORT_ERROR
else:
    try:
        if CHRONOS_VENDOR_SRC.exists() and str(CHRONOS_VENDOR_SRC) not in sys.path:
            sys.path.insert(0, str(CHRONOS_VENDOR_SRC))
        from chronos import ChronosPipeline as _ChronosPipeline  # type: ignore[import-not-found]

        ChronosPipeline = _ChronosPipeline
        CHRONOS_AVAILABLE = True
    except Exception as exc:  # pragma: no cover - import environment dependent
        CHRONOS_IMPORT_ERROR = str(exc)


class ChronosForecaster:
    key = "chronos"
    label = "Chronos"

    MODEL_HF_ID = os.getenv("CHRONOS_MODEL_HF_ID", "amazon/chronos-t5-tiny")
    MIN_CONTEXT = 32
    MAX_CONTEXT = 512
    DEFAULT_SAMPLE_COUNT = 32
    DEFAULT_TEMPERATURE = 1.0
    DEFAULT_TOP_K = 50
    DEFAULT_TOP_P = 1.0

    def __init__(self, logger: logging.Logger | None = None) -> None:
        self._logger = logger or logging.getLogger("ai-forecast.chronos")
        self._pipeline: Optional[Any] = None
        self._loaded = False
        self._load_lock: Optional[asyncio.Lock] = None
        self._predict_lock: Optional[asyncio.Lock] = None

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

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

    async def _get_load_lock(self) -> asyncio.Lock:
        if self._load_lock is None:
            self._load_lock = asyncio.Lock()
        return self._load_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._pipeline is None:
            return "not_loaded"
        try:
            return str(self._pipeline.model.device)
        except Exception:
            return "cpu"

    @staticmethod
    def _extract_timestamps(df: pd.DataFrame) -> pd.Series:
        if "timestamps" in df.columns:
            timestamps = pd.to_datetime(df["timestamps"], utc=True)
        elif "time" in df.columns:
            timestamps = pd.to_datetime(df["time"], unit="s", utc=True)
        else:
            raise HTTPException(status_code=422, detail="Chronos requires timestamps or time column")
        if timestamps.isna().any():
            raise HTTPException(status_code=422, detail="Chronos input timestamps contain NaT")
        return timestamps.reset_index(drop=True)

    @classmethod
    def _extract_ohlc4_series(cls, df: pd.DataFrame) -> np.ndarray:
        required_columns = {"open", "high", "low", "close"}
        missing_columns = sorted(required_columns - set(df.columns))
        if missing_columns:
            raise HTTPException(
                status_code=422,
                detail=f"Chronos input is missing OHLC columns: {', '.join(missing_columns)}",
            )

        ohlc4 = (
            df[["open", "high", "low", "close"]]
            .mean(axis=1)
            .to_numpy(dtype=np.float32, copy=False)
        )
        if ohlc4.ndim != 1:
            raise HTTPException(status_code=422, detail="Chronos input series must be 1-D")
        if len(ohlc4) < cls.MIN_CONTEXT:
            raise HTTPException(
                status_code=422,
                detail=f"Chronos requires at least {cls.MIN_CONTEXT} OHLC4 points",
            )
        if not np.isfinite(ohlc4).all():
            raise HTTPException(status_code=422, detail="Chronos input contains non-finite OHLC4 values")
        return ohlc4

    @staticmethod
    def _stable_seed(
        symbol: str,
        interval: str,
        horizon: int,
        context_length: int,
        last_timestamp: pd.Timestamp,
    ) -> int:
        payload = "|".join(
            [
                symbol,
                interval,
                str(horizon),
                str(context_length),
                str(int(last_timestamp.timestamp())),
            ]
        )
        return int(hashlib.sha256(payload.encode("utf-8")).hexdigest()[:8], 16)

    @staticmethod
    def _validate_output(
        quantiles: np.ndarray,
        mean_forecast: np.ndarray,
        horizon: int,
    ) -> Tuple[np.ndarray, np.ndarray, np.ndarray, Dict[str, Any]]:
        quantile_array = np.asarray(quantiles, dtype=float)
        mean_array = np.asarray(mean_forecast, dtype=float)

        if quantile_array.ndim != 3 or quantile_array.shape[0] != 1 or quantile_array.shape[1] < horizon or quantile_array.shape[2] < 3:
            raise HTTPException(
                status_code=500,
                detail=f"Chronos quantile forecast has invalid shape: {list(quantile_array.shape)}",
            )
        if mean_array.ndim != 2 or mean_array.shape[0] != 1 or mean_array.shape[1] < horizon:
            raise HTTPException(
                status_code=500,
                detail=f"Chronos mean forecast has invalid shape: {list(mean_array.shape)}",
            )

        p10 = quantile_array[0, :horizon, 0].astype(float)
        p50 = quantile_array[0, :horizon, 1].astype(float)
        p90 = quantile_array[0, :horizon, 2].astype(float)
        mean_values = mean_array[0, :horizon].astype(float)

        if not (np.isfinite(p10).all() and np.isfinite(p50).all() and np.isfinite(p90).all()):
            raise HTTPException(status_code=500, detail="Chronos 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="Chronos returned non-monotonic quantiles")

        return p10, p50, p90, {
            "quantile_shape": [int(dim) for dim in quantile_array.shape],
            "mean_shape": [int(dim) for dim in mean_array.shape],
            "quantiles_monotonic": quantiles_monotonic,
            "median_close_to_mean": bool(np.allclose(p50, mean_values, atol=1e-3, rtol=1e-3)),
        }

    async def _lazy_load(self) -> None:
        if self._loaded:
            return
        load_lock = await self._get_load_lock()
        async with load_lock:
            if self._loaded:
                return
            if not CHRONOS_AVAILABLE or ChronosPipeline is None:
                raise HTTPException(
                    status_code=503,
                    detail=(
                        "Chronos is unavailable. Ensure libs/chronos-forecasting is present "
                        "and its dependencies are installed."
                    ),
                )
            try:
                self._logger.info("[Chronos] Loading %s ...", self.MODEL_HF_ID)
                device_map = "cuda" if torch is not None and torch.cuda.is_available() else "cpu"
                model_dtype = torch.bfloat16 if torch is not None and torch.cuda.is_available() else torch.float32
                pipeline = await asyncio.to_thread(
                    ChronosPipeline.from_pretrained,
                    self.MODEL_HF_ID,
                    device_map=device_map,
                    dtype=model_dtype,
                )
                self._pipeline = pipeline
                self._loaded = True
                self._logger.info("[Chronos] Ready on %s", self.device)
            except Exception as exc:
                self._logger.error("[Chronos] Init failed: %s", exc, exc_info=True)
                raise HTTPException(status_code=500, detail=f"Chronos init failed: {exc}")

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

        timestamps = self._extract_timestamps(df)
        ohlc4_series = self._extract_ohlc4_series(df)
        context_length = min(len(ohlc4_series), self.MAX_CONTEXT, int(self._pipeline.model_context_length))
        context = ohlc4_series[-context_length:]
        seed = self._stable_seed(
            symbol=symbol or "UNKNOWN",
            interval=interval,
            horizon=horizon,
            context_length=context_length,
            last_timestamp=pd.Timestamp(timestamps.iloc[-1]).tz_convert("UTC"),
        )

        predict_lock = await self._get_predict_lock()
        async with predict_lock:
            try:
                random.seed(seed)
                np.random.seed(seed % (2**32 - 1))
                if torch is not None:
                    torch.manual_seed(seed)
                    if torch.cuda.is_available():
                        torch.cuda.manual_seed_all(seed)

                started_at = time.time()
                context_tensor = torch.tensor(context, dtype=torch.float32)
                quantiles, mean_forecast = await asyncio.to_thread(
                    self._pipeline.predict_quantiles,
                    context_tensor,
                    prediction_length=horizon,
                    quantile_levels=[0.1, 0.5, 0.9],
                    num_samples=self.DEFAULT_SAMPLE_COUNT,
                    temperature=self.DEFAULT_TEMPERATURE,
                    top_k=self.DEFAULT_TOP_K,
                    top_p=self.DEFAULT_TOP_P,
                    limit_prediction_length=False,
                )
                elapsed = time.time() - started_at
                self._logger.info(
                    "[Chronos] %.2fs | horizon=%d ctx=%d samples=%d device=%s",
                    elapsed,
                    horizon,
                    context_length,
                    self.DEFAULT_SAMPLE_COUNT,
                    self.device,
                )
            except Exception as exc:
                self._logger.error("[Chronos] Forecast failed: %s", exc, exc_info=True)
                raise HTTPException(status_code=500, detail=f"Chronos prediction failed: {exc}")

        p10, p50, p90, output_validation = self._validate_output(
            quantiles=quantiles.numpy() if hasattr(quantiles, "numpy") else np.asarray(quantiles),
            mean_forecast=mean_forecast.numpy() if hasattr(mean_forecast, "numpy") else np.asarray(mean_forecast),
            horizon=horizon,
        )

        return {
            "p10": p10,
            "p50": p50,
            "p90": p90,
            "model_name": self.MODEL_HF_ID,
            "context_length": context_length,
            "output_horizon": horizon,
            "sample_count": self.DEFAULT_SAMPLE_COUNT,
            "seed": seed,
            "input_semantics": {
                "feature_channels": ["ohlc4"],
                "active_forecast_channels": ["ohlc4"],
                "ignored_channels": ["volume", "amount"],
                "price_mode": "ohlc4_single_channel",
                "base_signal": "ohlc4",
                "volume_mode": "omitted",
                "amount_mode": "omitted",
                "adapter_mode": "chronos_univariate",
            },
            "input_validation": {
                "series_field": "ohlc4",
                "dtype": "float32",
                "context_length": context_length,
                "finite": True,
            },
            "output_semantics": {
                "forecast_channel": "ohlc4",
                "forecast_mode": "single_future_ohlc4_line",
                "quantile_fields": ["p10", "p50", "p90"],
                "quantile_source": "chronos_sampling",
                "candle_projection": "omitted",
                "reference_baseline": "last_ohlc4",
            },
            "output_validation": output_validation,
        }