service.py

from __future__ import annotations

import math
import os
from statistics import mean
from typing import Any

from schemas import HealthResponse, PredictRequest, PredictResponse, PredictionItem


class ChronosService:
    """HF Space service wrapper for Chronos-2."""

    def __init__(self) -> None:
        self.model_id = "chronos"
        self.model_name = os.getenv(
            "CHRONOS_MODEL_NAME",
            "amazon/chronos-2",
        )
        self.backend = os.getenv("CHRONOS_BACKEND", "hf_cpu").strip() or "hf_cpu"
        self.device = "cpu"
        self.max_context_length = int(os.getenv("CHRONOS_MAX_CONTEXT_LENGTH", "512"))
        self.max_horizon_step = int(os.getenv("CHRONOS_MAX_HORIZON_STEP", "288"))
        self.confidence_floor = float(os.getenv("CHRONOS_CONFIDENCE_FLOOR", "0.16"))
        self.confidence_ceiling = float(os.getenv("CHRONOS_CONFIDENCE_CEILING", "0.80"))
        self.min_required_points = int(os.getenv("CHRONOS_MIN_REQUIRED_POINTS", "32"))
        self.num_samples = int(os.getenv("CHRONOS_NUM_SAMPLES", "20"))
        self.allow_baseline_fallback = os.getenv("CHRONOS_ALLOW_BASELINE_FALLBACK", "false").lower() == "true"

        self.ready = False
        self.load_error = ""
        self._torch = None
        self._pipeline = None

        self._initialize_backend()

    def health(self) -> HealthResponse:
        return HealthResponse(
            status="ok" if self.ready else "degraded",
            model=self.model_name,
            model_id=self.model_id,
            backend=self.backend,
            device=self.device,
            ready=self.ready,
            max_context_length=self.max_context_length,
            max_horizon_step=self.max_horizon_step,
        )

    def predict(self, payload: PredictRequest) -> PredictResponse:
        self._validate_request(payload)
        closes = payload.close_prices[-payload.context_length :]

        if self.backend == "hf_cpu":
            if not self.ready:
                raise RuntimeError(self.load_error or "chronos backend not ready")
            predictions = self._predict_with_hf(closes, payload.horizons)
        else:
            predictions = self._predict_with_baseline(closes, payload.horizons)

        return PredictResponse(model_id=self.model_id, predictions=predictions)

    def _initialize_backend(self) -> None:
        if self.backend == "baseline_cpu":
            self.ready = True
            return
        if self.backend != "hf_cpu":
            raise ValueError(f"unsupported CHRONOS_BACKEND={self.backend}")

        try:
            self._load_hf_model()
            self.ready = True
        except Exception as exc:
            self.load_error = f"chronos hf load failed: {exc}"
            if self.allow_baseline_fallback:
                self.backend = "baseline_cpu"
                self.ready = True
            else:
                self.ready = False

    def _load_hf_model(self) -> None:
        import torch
        from chronos import Chronos2Pipeline

        self._torch = torch
        torch.set_num_threads(max(1, int(os.getenv("CHRONOS_TORCH_THREADS", "2"))))
        self._pipeline = Chronos2Pipeline.from_pretrained(
            self.model_name,
            device_map="cpu",
        )

    def _predict_with_hf(
        self, close_prices: list[float], horizons: list[int]
    ) -> list[PredictionItem]:
        assert self._torch is not None
        assert self._pipeline is not None

        torch = self._torch
        # Chronos-2 expects (n_series, n_variates, history_length) for tensor input.
        context = torch.tensor(
            close_prices[-self.max_context_length :],
            dtype=torch.float32,
        ).reshape(1, 1, -1)
        forecast = self._pipeline.predict(
            context,
            prediction_length=self.max_horizon_step,
        )

        dense_mean, dense_conf = self._extract_forecast(forecast)
        if len(dense_mean) < max(horizons):
            raise RuntimeError(
                f"Chronos output horizon {len(dense_mean)} is shorter than requested {max(horizons)}"
            )

        predictions: list[PredictionItem] = []
        for step in horizons:
            predictions.append(
                PredictionItem(
                    step=step,
                    pred_price=round(max(0.00000001, float(dense_mean[step - 1])), 8),
                    pred_confidence=round(dense_conf[step - 1], 4),
                )
            )
        return predictions

    def _extract_forecast(self, forecast: Any) -> tuple[list[float], list[float]]:
        assert self._torch is not None
        assert self._pipeline is not None
        torch = self._torch

        if isinstance(forecast, (list, tuple)):
            if not forecast:
                raise RuntimeError("empty Chronos forecast output")
            if len(forecast) != 1:
                raise RuntimeError(f"unexpected Chronos batch size: {len(forecast)}")
            forecast = forecast[0]

        if hasattr(forecast, "detach"):
            tensor = forecast.detach().cpu()
        else:
            tensor = torch.as_tensor(forecast)

        tensor = tensor.to(dtype=torch.float32)
        squeezed = tensor.squeeze()
        if squeezed.ndim == 0:
            raise RuntimeError(f"unexpected Chronos forecast shape: {tuple(tensor.shape)}")
        if squeezed.ndim == 1:
            mean_forecast = squeezed.tolist()
            std_forecast = [0.0 for _ in mean_forecast]
        else:
            if squeezed.ndim == 2:
                quantile_tensor = squeezed
            else:
                quantile_tensor = squeezed.reshape(-1, squeezed.shape[-2], squeezed.shape[-1]).mean(dim=0)

            quantiles = list(getattr(self._pipeline, "quantiles", []))
            if not quantiles:
                median_idx = quantile_tensor.shape[0] // 2
                mean_forecast = quantile_tensor[median_idx].tolist()
                std_forecast = [0.0 for _ in mean_forecast]
            else:
                median_idx = min(range(len(quantiles)), key=lambda idx: abs(quantiles[idx] - 0.5))
                lower_idx = min(range(len(quantiles)), key=lambda idx: abs(quantiles[idx] - 0.1))
                upper_idx = min(range(len(quantiles)), key=lambda idx: abs(quantiles[idx] - 0.9))
                mean_forecast = quantile_tensor[median_idx].tolist()
                lower_forecast = quantile_tensor[lower_idx].tolist()
                upper_forecast = quantile_tensor[upper_idx].tolist()
                std_forecast = [
                    max(0.0, (float(upper) - float(lower)) / 2.0)
                    for lower, upper in zip(lower_forecast, upper_forecast)
                ]

        confidence: list[float] = []
        for pred, std in zip(mean_forecast, std_forecast):
            dispersion = abs(float(std)) / max(abs(float(pred)), 1e-6)
            raw = 1.0 / (1.0 + dispersion)
            confidence.append(max(self.confidence_floor, min(self.confidence_ceiling, raw)))
        return [float(item) for item in mean_forecast], confidence

    def _validate_request(self, payload: PredictRequest) -> None:
        if payload.context_length > self.max_context_length:
            raise ValueError(
                f"context_length {payload.context_length} exceeds "
                f"CHRONOS_MAX_CONTEXT_LENGTH={self.max_context_length}"
            )
        if payload.context_length > len(payload.close_prices):
            raise ValueError("context_length must not exceed len(close_prices)")
        if len(payload.close_prices) < self.min_required_points:
            raise ValueError(
                f"at least {self.min_required_points} close prices are required "
                "for Chronos stability"
            )
        if any(step > self.max_horizon_step for step in payload.horizons):
            raise ValueError(
                f"horizons contain values above CHRONOS_MAX_HORIZON_STEP={self.max_horizon_step}"
            )

    def _predict_with_baseline(
        self, close_prices: list[float], horizons: list[int]
    ) -> list[PredictionItem]:
        last_price = close_prices[-1]
        short_window = close_prices[-min(10, len(close_prices)) :]
        mid_window = close_prices[-min(24, len(close_prices)) :]
        long_window = close_prices[-min(64, len(close_prices)) :]

        short_mean = mean(short_window)
        mid_mean = mean(mid_window)
        long_mean = mean(long_window)
        momentum = 0.0 if short_mean == 0 else (last_price - short_mean) / short_mean
        mean_reversion = 0.0 if long_mean == 0 else (mid_mean - long_mean) / long_mean
        local_trend = self._slope(mid_window)

        predictions: list[PredictionItem] = []
        for step in horizons:
            horizon_scale = min(1.0, math.log(step + 1.0) / 3.8)
            expected_return = momentum * 0.35 + mean_reversion * 0.30 + local_trend * 0.35
            expected_return *= horizon_scale

            pred_price = max(0.00000001, last_price * (1.0 + expected_return))
            confidence = self._baseline_confidence(close_prices, step, abs(expected_return))
            predictions.append(
                PredictionItem(
                    step=step,
                    pred_price=round(pred_price, 8),
                    pred_confidence=round(confidence, 4),
                )
            )
        return predictions

    def _baseline_confidence(
        self, close_prices: list[float], step: int, expected_move_abs: float
    ) -> float:
        if len(close_prices) < 3:
            return self.confidence_floor

        changes: list[float] = []
        for previous, current in zip(close_prices[:-1], close_prices[1:]):
            if previous <= 0:
                continue
            changes.append(abs((current - previous) / previous))

        realized_vol = mean(changes[-min(64, len(changes)) :]) if changes else 0.0
        stability = max(0.0, 1.0 - min(realized_vol * 18.0, 1.0))
        horizon_decay = 1.0 / (1.0 + math.log(step + 1.0))
        raw = 0.20 + min(expected_move_abs / (realized_vol + 1e-9), 2.0) * 0.18
        raw += stability * 0.20 + horizon_decay * 0.22
        return max(self.confidence_floor, min(self.confidence_ceiling, raw))

    @staticmethod
    def _slope(values: list[float]) -> float:
        if len(values) < 2 or values[0] == 0:
            return 0.0
        return (values[-1] - values[0]) / values[0]

    def describe_runtime(self) -> dict[str, Any]:
        return {
            "model_id": self.model_id,
            "model_name": self.model_name,
            "backend": self.backend,
            "device": self.device,
            "ready": self.ready,
            "load_error": self.load_error,
            "max_context_length": self.max_context_length,
            "max_horizon_step": self.max_horizon_step,
            "min_required_points": self.min_required_points,
        }