from __future__ import annotations
import numpy as np
from typing import Iterator, List, Optional
from gluonts.dataset import Dataset as GluonDataset
from gluonts.model import Forecast
from gluonts.model.forecast import QuantileForecast
from gluonts.model.predictor import RepresentablePredictor
from tsfm_bench.eval.predictors import _ForecastConfig

class MyUserModel(RepresentablePredictor):
    """A mock user-submitted pre-trained TSFM (e.g., an autoregressive or slightly perturbed model)."""

    def __init__(
        self,
        prediction_length: int,
        checkpoint_path: str | None = None,
        quantile_levels: Optional[List[float]] = None,
    ):
        super().__init__(prediction_length=prediction_length)
        self.config = _ForecastConfig.from_quantiles(quantile_levels)
        self.checkpoint_path = checkpoint_path
        # In a real model, we would load model weights here:
        # self.model = load_model_from_checkpoint(checkpoint_path)

    def predict(self, dataset: GluonDataset, **kwargs) -> Iterator[Forecast]:
        for entry in dataset:
            target = np.asarray(entry["target"], dtype=np.float64)
            univariate = target.ndim == 1
            if univariate:
                working = target.reshape(1, -1)
            else:
                working = target

            # Autoregressive/perturbed simulation:
            # We take the last few points and project them forward with a small dampening factor
            # and a tiny bit of deterministic shift (to make it different from SeasonalNaive)
            horizon = self.prediction_length
            last_val = working[:, -1:]

            # Simple linear trend projection with noise dampening
            predictions = []
            for h in range(1, horizon + 1):
                val = last_val * (0.985 ** h) + (np.sin(h) * 0.005 * last_val)
                predictions.append(val)
            mean = np.concatenate(predictions, axis=1)

            if univariate:
                mean = mean[0]

            arrays = [mean]
            if self.config.quantile_levels:
                for q in self.config.quantile_levels:
                    # Deterministic quantiles for simulation
                    if q <= 0.5:
                        arrays.append(mean * (1 - (0.5 - q) * 0.08))
                    else:
                        arrays.append(mean * (1 + (q - 0.5) * 0.08))

            forecast_arrays = np.stack(arrays, axis=0)

            yield QuantileForecast(
                forecast_arrays=forecast_arrays,
                forecast_keys=self.config.forecast_keys,
                start_date=entry["start"],
                item_id=entry["item_id"],
            )