models.py

"""Model wrappers for TimesFM and Chronos-2 with graceful degradation.

Both wrappers expose a uniform interface::

    model.available -> bool
    model.label     -> str   (e.g. "TimesFM 2.5", "Chronos-2")
    model.forecast(series, horizon, covariates=None) -> np.ndarray

If the heavy ML dependencies or the model weights are unavailable (not
installed, offline, download failed, OOM, ...) the wrapper keeps
``available = False`` and ``forecast`` returns a damped naive forecast.
That guarantees ``main.py`` always starts and ``/forecast`` always responds,
which is exactly what Phase 5 asks for. When the real weights are present the
wrappers use them.

The API surfaces of these libraries changed across releases, so each loader
tries the current (2.x) API first and falls back to the legacy API. Every
inference call is wrapped so a runtime surprise degrades to the naive path
instead of 500-ing the request.
"""

from __future__ import annotations

import logging
from typing import Optional, Sequence

import numpy as np

log = logging.getLogger("forecast.models")


def naive_forecast(series: Sequence[float], horizon: int) -> np.ndarray:
    """Damped-drift fallback forecast.

    Continues the last observed step, damping it each period so the line does
    not run away over the horizon. Falls back to persistence (repeat the last
    value) when there is too little history.
    """
    arr = np.asarray(series, dtype=float)
    arr = arr[np.isfinite(arr)]
    if arr.size == 0:
        return np.zeros(int(horizon), dtype=float)

    last = float(arr[-1])
    step = float(arr[-1] - arr[-2]) if arr.size >= 2 else 0.0
    damp = 0.6
    out = []
    cur = last
    for _ in range(int(horizon)):
        cur = cur + step
        step *= damp
        out.append(cur)
    return np.asarray(out, dtype=float)


class TimesFMModel:
    """Google TimesFM wrapper (tries the 2.5 API, then the legacy 2.0 API)."""

    def __init__(self) -> None:
        self.available = False
        self.label = "TimesFM 2.0"
        self._impl = None
        self._mode: Optional[str] = None
        self._load()

    def _load(self) -> None:
        try:
            import timesfm  # type: ignore
        except Exception as exc:  # pragma: no cover - import guard
            log.warning("timesfm not importable (%s) — using naive fallback", exc)
            return

        # --- Preferred: TimesFM 2.5 (current PyPI API) ---------------------
        try:
            import torch  # type: ignore

            torch.set_float32_matmul_precision("high")
            model = timesfm.TimesFM_2p5_200M_torch.from_pretrained(
                "google/timesfm-2.5-200m-pytorch"
            )
            model.compile(
                timesfm.ForecastConfig(
                    max_context=1024,
                    max_horizon=256,
                    normalize_inputs=True,
                    use_continuous_quantile_head=True,
                    force_flip_invariance=True,
                    infer_is_positive=False,  # probabilities can sit near 0
                    fix_quantile_crossing=True,
                )
            )
            self._impl = model
            self._mode = "2.5"
            self.label = "TimesFM 2.5"
            self.available = True
            log.info("Loaded TimesFM 2.5 (google/timesfm-2.5-200m-pytorch)")
            return
        except Exception as exc:
            log.warning("TimesFM 2.5 load failed (%s) — trying legacy 2.0 API", exc)

        # --- Fallback: TimesFM 2.0 (legacy hparams/checkpoint API) ---------
        try:
            tfm = timesfm.TimesFm(
                hparams=timesfm.TimesFmHparams(
                    backend="cpu",
                    per_core_batch_size=32,
                    horizon_len=128,
                    context_len=2048,
                ),
                checkpoint=timesfm.TimesFmCheckpoint(
                    huggingface_repo_id="google/timesfm-2.0-500m-pytorch"
                ),
            )
            self._impl = tfm
            self._mode = "2.0"
            self.label = "TimesFM 2.0"
            self.available = True
            log.info("Loaded TimesFM 2.0 (google/timesfm-2.0-500m-pytorch)")
        except Exception as exc:
            log.warning("TimesFM 2.0 load failed (%s) — using naive fallback", exc)

    def forecast(
        self,
        series: Sequence[float],
        horizon: int,
        covariates: Optional[Sequence[Sequence[float]]] = None,
    ) -> np.ndarray:
        if not self.available or self._impl is None:
            return naive_forecast(series, horizon)

        arr = np.asarray(series, dtype=float)
        try:
            if self._mode == "2.5":
                point, _quantiles = self._impl.forecast(horizon=int(horizon), inputs=[arr])
                return np.asarray(point[0], dtype=float)[:horizon]
            # legacy 2.0
            point, _ = self._impl.forecast([arr], freq=[0])
            return np.asarray(point[0], dtype=float)[:horizon]
        except Exception as exc:
            log.warning("TimesFM forecast failed (%s) — naive fallback", exc)
            return naive_forecast(series, horizon)


class ChronosModel:
    """Amazon Chronos-2 wrapper (falls back to a Chronos-Bolt pipeline)."""

    def __init__(self) -> None:
        self.available = False
        self.label = "Chronos-2"
        self._impl = None
        self._mode: Optional[str] = None
        self._load()

    def _load(self) -> None:
        # --- Preferred: Chronos-2 -----------------------------------------
        try:
            from chronos import Chronos2Pipeline  # type: ignore

            self._impl = Chronos2Pipeline.from_pretrained(
                "amazon/chronos-2", device_map="cpu"
            )
            self._mode = "chronos2"
            self.label = "Chronos-2"
            self.available = True
            log.info("Loaded Chronos-2 (amazon/chronos-2)")
            return
        except Exception as exc:
            log.warning("Chronos-2 load failed (%s) — trying Chronos-Bolt", exc)

        # --- Fallback: Chronos-Bolt (older base pipeline) -----------------
        try:
            from chronos import BaseChronosPipeline  # type: ignore

            self._impl = BaseChronosPipeline.from_pretrained(
                "amazon/chronos-bolt-base", device_map="cpu"
            )
            self._mode = "bolt"
            self.label = "Chronos-Bolt"
            self.available = True
            log.info("Loaded Chronos-Bolt (amazon/chronos-bolt-base)")
        except Exception as exc:
            log.warning("Chronos-Bolt load failed (%s) — using naive fallback", exc)

    def forecast(
        self,
        series: Sequence[float],
        horizon: int,
        covariates: Optional[Sequence[Sequence[float]]] = None,
    ) -> np.ndarray:
        if not self.available or self._impl is None:
            return naive_forecast(series, horizon)

        arr = np.asarray(series, dtype=float)
        try:
            if self._mode == "chronos2":
                return self._forecast_chronos2(arr, horizon, covariates)
            return self._forecast_bolt(arr, horizon)
        except Exception as exc:
            log.warning("Chronos forecast failed (%s) — naive fallback", exc)
            return naive_forecast(series, horizon)

    def _forecast_chronos2(
        self,
        arr: np.ndarray,
        horizon: int,
        covariates: Optional[Sequence[Sequence[float]]],
    ) -> np.ndarray:
        import pandas as pd

        ts = pd.date_range("2000-01-01", periods=arr.size, freq="D")
        frame = {"id": ["series"] * arr.size, "timestamp": ts, "target": arr}

        # Chronos-2 natively supports past covariates; attach any we were given.
        if covariates:
            for idx, cov in enumerate(covariates):
                cov_arr = np.asarray(cov, dtype=float)
                if cov_arr.size == arr.size:
                    frame[f"cov_{idx}"] = cov_arr

        context_df = pd.DataFrame(frame)
        pred_df = self._impl.predict_df(
            context_df,
            prediction_length=int(horizon),
            quantile_levels=[0.1, 0.5, 0.9],
            id_column="id",
            timestamp_column="timestamp",
            target="target",
        )
        return self._extract_median(pred_df, horizon)

    @staticmethod
    def _extract_median(pred_df, horizon: int) -> np.ndarray:
        import pandas as pd  # noqa: F401

        # predict_df output column naming varies; prefer the 0.5 quantile/mean.
        for candidate in ("0.5", 0.5, "median", "mean", "predictions", "target"):
            if candidate in pred_df.columns:
                return np.asarray(pred_df[candidate].to_numpy(), dtype=float)[:horizon]
        # Last resort: first purely-numeric column that is not an id/time column.
        for col in pred_df.columns:
            if col in ("id", "timestamp", "item_id"):
                continue
            try:
                return np.asarray(pred_df[col].to_numpy(), dtype=float)[:horizon]
            except Exception:
                continue
        raise ValueError("could not locate a forecast column in predict_df output")

    def _forecast_bolt(self, arr: np.ndarray, horizon: int) -> np.ndarray:
        import torch  # type: ignore

        context = torch.tensor(arr, dtype=torch.float32)
        # BaseChronosPipeline exposes predict_quantiles(...) -> (quantiles, mean)
        quantiles, mean = self._impl.predict_quantiles(
            context=context,
            prediction_length=int(horizon),
            quantile_levels=[0.1, 0.5, 0.9],
        )
        median = quantiles[0, :, 1]  # the 0.5 quantile
        return np.asarray(median, dtype=float)[:horizon]