"""
Salesforce Moirai — a pretrained time-series foundation model, via the
`uni2ts` package (https://github.com/SalesforceAIResearch/uni2ts).

Needs: pip install uni2ts torch
First use downloads pretrained weights from Hugging Face (needs internet),
cached under ~/.cache/huggingface/ afterwards.

Uni2TS uses GluonTS for inference. The pattern below (PandasDataset ->
MoiraiForecast -> create_predictor -> predict) matches the official
"Getting Started" example in the uni2ts README.

Optional 10-feature covariate input (see features/feature_pipeline.py): of
all 5 models, Moirai has the most natively-built-for-this mechanism —
`past_feat_dynamic_real`, a first-class GluonTS/uni2ts concept for exactly
"extra real-valued series that are only observed for the historical
context, not the forecast horizon" (see
github.com/SalesforceAIResearch/uni2ts discussion #166 for the underlying
forward() signature). RSI/MACD/Bollinger/SMA-EMA/Stochastic/Volume for
candles that haven't happened yet are precisely that: unknown-until-they-
happen. Concretely, `PandasDataset(..., past_feat_dynamic_real=[...column
names...])` takes those columns straight from the same per-series
DataFrame that holds the target, and `MoiraiForecast(...,
past_feat_dynamic_real_dim=N)` tells the pretrained module how many such
columns to expect — both are genuine constructor parameters this project
already left wired up at `0` (see the two _dim kwargs below), not
something bolted on. Unlike the ARIMA family or TimesFM, this means Moirai
never needs carry_forward_future() or any other future-covariate stand-in
— it was designed to just not need future values for these at all.
"""
import numpy as np
import pandas as pd

from models.base_model import BaseForecastModel
from config import MOIRAI_CHECKPOINT
from features.feature_pipeline import validate_features


class MoiraiModel(BaseForecastModel):
    name = "Moirai"

    def __init__(self, checkpoint: str = MOIRAI_CHECKPOINT, context_length: int = 200,
                 num_samples: int = 100):
        self.checkpoint = checkpoint
        self.context_length = context_length
        self.num_samples = num_samples
        self._module = None  # cached after first load so repeat calls in the
        # same running app don't re-download / re-init the pretrained weights

    def _check_dependencies(self):
        import uni2ts  # noqa
        import gluonts  # noqa

    def _get_module(self):
        from uni2ts.model.moirai import MoiraiModule
        if self._module is None:
            self._module = MoiraiModule.from_pretrained(self.checkpoint)
        return self._module

    def predict(self, history: pd.Series, horizon: int = 1, features: pd.DataFrame = None) -> list:
        try:
            from uni2ts.model.moirai import MoiraiForecast
            from gluonts.dataset.pandas import PandasDataset
        except ImportError as e:
            raise ImportError(
                "Moirai needs the 'uni2ts' and 'gluonts' packages. Install with:\n"
                "  pip install uni2ts torch\n"
                f"(original error: {e})"
            )

        values = np.asarray(history, dtype=np.float32)
        if len(values) < 10:
            raise ValueError("Moirai needs at least 10 candles of history.")

        if features is not None:
            # Length-only here (against the FULL history) -- catches a
            # genuinely misaligned caller early. NaN is checked below,
            # only for the truncated tail actually used (see
            # validate_features' docstring for why the split matters).
            validate_features(features, expected_len=len(values), model_name="Moirai", check_nan=False)

        context = values[-self.context_length:] if len(values) > self.context_length else values

        # Moirai/GluonTS need a DatetimeIndex but only the VALUES matter for a
        # point forecast — the real candle timestamps are re-attached in
        # app.py, so a clean synthetic per-minute index is fine here.
        # (Timestamp.now("UTC"), not the deprecated Timestamp.utcnow() —
        # pandas already warns that utcnow() will be removed in a future
        # release.)
        idx = pd.date_range(end=pd.Timestamp.now("UTC").floor("min"), periods=len(context), freq="min")
        data = {"target": context}

        feat_cols = []
        if features is not None:
            # Same tail-truncation as `context` above, so every column in
            # `data` still describes the same set of candles, row for row.
            feat_context = features.iloc[-len(context):].reset_index(drop=True)
            validate_features(feat_context, expected_len=len(context), model_name="Moirai")
            for col in feat_context.columns:
                data[col] = feat_context[col].to_numpy(dtype=np.float32)
            feat_cols = list(feat_context.columns)

        series_df = pd.DataFrame(data, index=idx)
        dataset = PandasDataset(
            series_df, target="target",
            past_feat_dynamic_real=feat_cols if feat_cols else None,
        )

        module = self._get_module()
        forecaster = MoiraiForecast(
            module=module,
            prediction_length=horizon,
            context_length=len(context),
            patch_size="auto",
            num_samples=self.num_samples,
            target_dim=1,
            feat_dynamic_real_dim=0,
            past_feat_dynamic_real_dim=len(feat_cols),
        )
        predictor = forecaster.create_predictor(batch_size=1)
        forecast = next(iter(predictor.predict(dataset)))

        # forecast.samples has shape (num_samples, horizon) — take the
        # per-step median across samples as the point forecast.
        point_forecast = np.median(forecast.samples, axis=0)
        out = [float(x) for x in point_forecast]
        if not all(np.isfinite(x) for x in out):
            raise ValueError("Moirai produced a non-finite forecast (NaN/inf) for this window.")
        return out