import numpy as np
import pandas as pd
import threading
from models.base_model import BaseForecastModel, validate_history_and_horizon
from config import MOIRAI_CHECKPOINT, RANDOM_SEED
from features.feature_pipeline import validate_features
_RNG_LOCK = threading.Lock()
_MODULE_CACHE = {}
_MODULE_LOCK = threading.Lock()


def _step_from_history(history: pd.Series) -> pd.Timedelta:
    if isinstance(history.index, pd.DatetimeIndex) and len(history.index) >= 2:
        diffs = history.index.to_series().diff().dropna()
        if not diffs.empty:
            step = diffs.median()
            if step > pd.Timedelta(0):
                if not diffs.eq(step).all():
                    raise ValueError('Moirai needs a continuous history; irregular timestamps cannot be relabelled.')
                return step
    return pd.Timedelta(minutes=1)


class MoiraiModel(BaseForecastModel):
    name = "Moirai"

    def __init__(
        self,
        checkpoint: str = MOIRAI_CHECKPOINT,
        context_length: int = 256,
        num_samples: int = 20,
    ):
        self.checkpoint = checkpoint
        self.context_length = context_length
        self.num_samples = num_samples
        self._module = None
        self.last_diagnostics = {}

    def _get_module(self):
        from uni2ts.model.moirai import MoiraiModule

        if self._module is None:
            with _MODULE_LOCK:
                self._module = _MODULE_CACHE.get(self.checkpoint)
                if self._module is None:
                    self._module = MoiraiModule.from_pretrained(self.checkpoint)
                    self._module.eval()
                    _MODULE_CACHE[self.checkpoint] = self._module
        return self._module

    def predict(
        self, history: pd.Series, horizon: int = 1, features: pd.DataFrame = None
    ) -> list:
        self.last_diagnostics = {}
        validate_history_and_horizon(history, horizon, self.name)
        try:
            import torch
        except ImportError as error:
            raise ImportError('Moirai requires the optional foundation-model packages. Install requirements-foundation.txt.') from error
        with _RNG_LOCK, torch.random.fork_rng(devices=[]), torch.inference_mode():
            torch.manual_seed(RANDOM_SEED)
            return self._predict(history, horizon, features)

    def _predict(self, history, horizon, features):
        values = np.asarray(history, dtype=np.float32)
        if len(values) < 10:
            raise ValueError("Moirai needs at least 10 candles of history.")
        validate_history_and_horizon(values, horizon, self.name)
        try:
            from uni2ts.model.moirai import MoiraiForecast
            from gluonts.dataset.pandas import PandasDataset
        except ImportError as e:
            raise ImportError(
                f"Moirai needs the 'uni2ts' and 'gluonts' packages. Install with:\n  pip install uni2ts torch\n(original error: {e})"
            )
        if features is not None:
            validate_features(
                features,
                expected_len=len(values),
                model_name="Moirai",
                check_nan=False,
            )
        # Bound the context for the small checkpoint so a large dashboard
        # window cannot cause an avoidable CPU-memory spike.
        context_length = min(max(int(self.context_length), 32), 512)
        context = values[-context_length:] if len(values) > context_length else values
        step = _step_from_history(history)
        end_ts = pd.Timestamp(history.index[-1]) if isinstance(history.index, pd.DatetimeIndex) else pd.Timestamp('2000-01-01')
        if end_ts.tzinfo is not None:
            end_ts = end_ts.tz_convert('UTC').tz_localize(None)
        idx = pd.date_range(end=end_ts, periods=len(context), freq=step)
        data = {"target": context}
        feat_cols = []
        if features is not None:
            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)))
        samples = np.asarray(forecast.samples, dtype=float)
        if samples.ndim == 3 and samples.shape[-1] == 1:
            samples = samples[..., 0]
        if samples.ndim != 2 or samples.shape != (self.num_samples, horizon):
            raise ValueError('Moirai returned an unexpected univariate sample/horizon shape.')
        out = np.median(samples, axis=0).tolist()
        if not all(np.isfinite(x) and x > 0 for x in out):
            raise ValueError(
                "Moirai produced a non-finite forecast (NaN/inf) for this window."
            )
        self.last_diagnostics = {'checkpoint': self.checkpoint, 'context_candles': len(context), 'point_statistic': 'sample median', 'seed': RANDOM_SEED, 'fallback': False, 'weights_loaded': True}
        return out