# space/tools/ts_forecast_tool.py
import numpy as np
import pandas as pd
from typing import Dict, Optional

from utils.tracing import Tracer
from utils.config import AppConfig

# Granite TTM imports (from tsfm_public)
from tsfm_public.models.tinytimemixer.modeling_tinytimemixer import TinyTimeMixerForPrediction
from tsfm_public.toolkit.config import TSPPConfig
from tsfm_public.toolkit.dataset import TimeSeriesDataset
from tsfm_public.toolkit.trainer import Trainer


class TimeseriesForecastTool:
    """
    Wraps IBM Granite TinyTimeMixer (TTM) for multivariate forecasting.
    Expects a wide dataframe with:
      - 'timestamp' (datetime64[ns])
      - one or more numeric series columns to forecast (targets)
      - optional control/exogenous columns (known-in-future features)
    You must provide context_length and forecast_length to match a TTM variant.
    """
    def __init__(self, cfg: AppConfig, tracer: Tracer,
                 hf_model_id: str = "ibm-granite/granite-timeseries-ttm-r1",
                 revision: str = "main",
                 context_length: int = 512,
                 forecast_length: int = 96,
                 target_cols: Optional[list] = None,
                 control_cols: Optional[list] = None):
        self.cfg = cfg
        self.tracer = tracer

        self.context_length = context_length
        self.forecast_length = forecast_length
        self.target_cols = target_cols or []
        self.control_cols = control_cols or []

        # Build TSPP config
        self.tspp_config = TSPPConfig(
            context_length=context_length,
            prediction_length=forecast_length,
            target_cols=self.target_cols,
            known_cov_cols=self.control_cols,  # known-in-future exogenous
            time_col="timestamp",
            freq=None  # inferred; you can set "H" or "T" if you know it
        )

        # Load model from HF (r1; try r2 for newer variants if needed)
        self.model = TinyTimeMixerForPrediction.from_pretrained(
            hf_model_id, revision=revision
        )

    def _build_dataset(self, df: pd.DataFrame) -> TimeSeriesDataset:
        # Minimal build: single item dataset from dataframe (you can batch multiple series)
        item = df.sort_values("timestamp").reset_index(drop=True)
        return TimeSeriesDataset.from_pandas(
            item,
            tspp_config=self.tspp_config
        )

    def zeroshot_forecast(self, df: pd.DataFrame) -> Dict[str, pd.DataFrame]:
        """
        Zero-shot forecast: no fine-tuning, just apply pre-trained model.
        Returns dict with "forecast" (future horizon) and "context" (last context window).
        """
        dset = self._build_dataset(df)
        trainer = Trainer(model=self.model)
        out = trainer.evaluate(dset)  # Granite API uses 'evaluate' for zeroshot

        # Convert to a user-friendly dataframe
        # out contains predictions for target_cols for next forecast_length steps
        preds = out["predictions"]  # dict: {col: np.ndarray [forecast_length]}
        horizon_idx = pd.RangeIndex(self.forecast_length, name="step")
        forecast_df = pd.DataFrame(preds, index=horizon_idx).reset_index(drop=True)

        try:
            self.tracer.trace_event("ts_forecast", {
                "targets": self.target_cols,
                "ctx": self.context_length,
                "h": self.forecast_length
            })
        except Exception:
            pass

        return {
            "forecast": forecast_df,
            "context": df.tail(self.context_length).reset_index(drop=True)
        }