"""
Chronos-2 (amazon/chronos-2) integration.

Chronos-2 is a 120M-parameter, encoder-only zero-shot time series
foundation model from Amazon Science (released 20 Oct 2025), successor
to the original Chronos/Chronos-Bolt family. Unlike those (and unlike
Moirai-MoE), Chronos-2 produces quantile forecasts DIRECTLY rather than
by sampling paths — confirmed by reading chronos/chronos2/pipeline.py's
own source (`forecast_type: ForecastType = ForecastType.QUANTILES`) in
a network-enabled sandbox pass, not assumed from documentation. That
means Chronos-2's output for identical input is deterministic by
construction — it does not need the seed fix moirai_model.py required.

API confirmed two ways in that same sandbox pass: (1) the officially
documented usage pattern (huggingface.co/amazon/chronos-2, the
amazon-science/chronos-forecasting README), and (2) by installing
`chronos-forecasting>=2.0` with `--no-deps` (skipping torch, which
isn't installed here) and reading chronos/chronos2/pipeline.py's real
`predict_df()` signature, docstring, and implementation directly —
including its exact output column schema, which the README's example
doesn't fully spell out:

    from chronos import Chronos2Pipeline
    pipeline = Chronos2Pipeline.from_pretrained("amazon/chronos-2", device_map="cpu")
    pred_df = pipeline.predict_df(df, id_column=..., timestamp_column=...,
                                   target=..., prediction_length=H,
                                   quantile_levels=[0.1, 0.25, 0.5, 0.75, 0.9])
    # columns: id_column, timestamp_column, "target_name", "predictions"
    # (mean), and one column per quantile level named as str(level).

`chronos/__init__.py` imports torch unconditionally at module level
(confirmed directly — unlike uni2ts's lazier import style), so even
constructing this class must defer the import to load()/predict(),
same reasoning as moirai_model.py's lazy imports but enforced one
level higher up.

NETWORK NOTE: load() downloads the real checkpoint from Hugging Face on
first use, same as moirai_model.py. In this project's sandbox passes so
far, that download has never succeeded (no route to huggingface.co) --
confirmed as ModelLoadError with the underlying exception attached, not
a fabricated result, matching moirai_model.py's own behavior in the
same situation.
"""
from __future__ import annotations

import numpy as np
import pandas as pd

from forecast_types import ModelInferenceError, ModelInputError, ModelLoadError, QuantileForecast

CHECKPOINT = "amazon/chronos-2"
QUANTILE_LEVELS = [0.1, 0.25, 0.5, 0.75, 0.9]


class Chronos2Forecaster:
    """Same load()/predict() contract as MoiraiMoEForecaster and
    AutoArimaForecaster (see forecast_types.py) — no shared code with
    either, no fallback to either."""

    def __init__(self, checkpoint: str = CHECKPOINT, device: str = "cpu", min_context: int = 32):
        self.checkpoint = checkpoint
        self.device = device
        self.min_context = min_context
        self.max_context = 2048  # Chronos2Pipeline.default_context_length, confirmed from source
        self._pipeline = None
        self._loaded = False

    def load(self):
        try:
            from chronos import Chronos2Pipeline
        except ImportError as e:
            raise ModelLoadError(
                f"chronos-forecasting (with torch) is not installed: {e}"
            ) from e
        try:
            self._pipeline = Chronos2Pipeline.from_pretrained(self.checkpoint, device_map=self.device)
        except Exception as e:
            raise ModelLoadError(f"failed to load {self.checkpoint}: {e}") from e
        self._loaded = True

    def predict(self, context_df: pd.DataFrame, target_col: str, past_covariate_cols,
                prediction_length: int, freq: str) -> QuantileForecast:
        if not self._loaded:
            raise ModelInferenceError("call load() before predict()")

        closes = context_df[target_col]
        if closes.isna().any() or not np.isfinite(closes.to_numpy()).all():
            raise ModelInputError("context contains NaN/inf values in the target column")

        context_length = min(len(context_df), self.max_context)
        if context_length < self.min_context:
            raise ModelInputError(
                f"context too short: {context_length} row(s), need at least {self.min_context}"
            )
        window = context_df.iloc[-context_length:]

        in_df = pd.DataFrame({
            "item_id": "series",
            "timestamp": window.index,
            "target": window[target_col].to_numpy(),
        })
        try:
            pred_df = self._pipeline.predict_df(
                in_df, id_column="item_id", timestamp_column="timestamp", target="target",
                prediction_length=prediction_length, quantile_levels=QUANTILE_LEVELS,
            )
        except Exception as e:
            raise ModelInferenceError(f"Chronos-2 predict_df failed: {e}") from e

        try:
            quantiles = pred_df[[str(level) for level in QUANTILE_LEVELS]].to_numpy()
        except KeyError as e:
            raise ModelInferenceError(
                f"predict_df's output columns didn't match the expected quantile-level "
                f"names {[str(l) for l in QUANTILE_LEVELS]} -- got {list(pred_df.columns)}. "
                f"The installed chronos-forecasting version's output schema may have changed "
                f"since this was last verified against source (spec section 75)."
            ) from e

        return QuantileForecast(
            horizon=prediction_length, quantile_levels=QUANTILE_LEVELS,
            quantiles=quantiles, samples=None,
            context_length_used=context_length, prediction_length=prediction_length,
        )