src/gift_eval/predictor.py

"""Predictor implementation wrapping the TimeSeriesModel for GIFT-Eval."""

import logging
from collections.abc import Iterator

import numpy as np
import torch
import yaml
from gluonts.model.forecast import QuantileForecast
from gluonts.model.predictor import Predictor
from torch.nn.parallel import DistributedDataParallel as DDP

from src.data.containers import BatchTimeSeriesContainer
from src.data.frequency import parse_frequency
from src.data.scalers import RobustScaler
from src.models.model import TimeSeriesModel
from src.utils.utils import device

logger = logging.getLogger(__name__)


class TimeSeriesPredictor(Predictor):
    """Unified predictor for TimeSeriesModel supporting flexible construction."""

    def __init__(
        self,
        model: TimeSeriesModel,
        config: dict,
        ds_prediction_length: int,
        ds_freq: str,
        batch_size: int = 32,
        max_context_length: int | None = None,
        debug: bool = False,
    ) -> None:
        # Dataset-specific context (can be updated per dataset/term)
        self.ds_prediction_length = ds_prediction_length
        self.ds_freq = ds_freq
        self.batch_size = batch_size
        self.max_context_length = max_context_length
        self.debug = debug

        # Persistent model/config (unwrap DDP if needed)
        self.model = model.module if isinstance(model, DDP) else model
        self.model.eval()
        self.config = config

        # Initialize scaler (using same type as model)
        scaler_type = self.config.get("TimeSeriesModel", {}).get("scaler", "custom_robust")
        epsilon = self.config.get("TimeSeriesModel", {}).get("epsilon", 1e-3)
        if scaler_type == "custom_robust":
            self.scaler = RobustScaler(epsilon=epsilon)
        else:
            raise ValueError(f"Unsupported scaler type: {scaler_type}")

    def set_dataset_context(
        self,
        prediction_length: int | None = None,
        freq: str | None = None,
        batch_size: int | None = None,
        max_context_length: int | None = None,
    ) -> None:
        """Update lightweight dataset-specific attributes without reloading the model."""

        if prediction_length is not None:
            self.ds_prediction_length = prediction_length
        if freq is not None:
            self.ds_freq = freq
        if batch_size is not None:
            self.batch_size = batch_size
        if max_context_length is not None:
            self.max_context_length = max_context_length

    @classmethod
    def from_model(
        cls,
        model: TimeSeriesModel,
        config: dict,
        ds_prediction_length: int,
        ds_freq: str,
        batch_size: int = 32,
        max_context_length: int | None = None,
        debug: bool = False,
    ) -> "TimeSeriesPredictor":
        return cls(
            model=model,
            config=config,
            ds_prediction_length=ds_prediction_length,
            ds_freq=ds_freq,
            batch_size=batch_size,
            max_context_length=max_context_length,
            debug=debug,
        )

    @classmethod
    def from_paths(
        cls,
        model_path: str,
        config_path: str,
        ds_prediction_length: int,
        ds_freq: str,
        batch_size: int = 32,
        max_context_length: int | None = None,
        debug: bool = False,
    ) -> "TimeSeriesPredictor":
        with open(config_path) as f:
            config = yaml.safe_load(f)
        model = cls._load_model_from_path(config=config, model_path=model_path)
        return cls(
            model=model,
            config=config,
            ds_prediction_length=ds_prediction_length,
            ds_freq=ds_freq,
            batch_size=batch_size,
            max_context_length=max_context_length,
            debug=debug,
        )

    @staticmethod
    def _load_model_from_path(config: dict, model_path: str) -> TimeSeriesModel:
        try:
            model = TimeSeriesModel(**config["TimeSeriesModel"]).to(device)
            checkpoint = torch.load(model_path, map_location=device)
            model.load_state_dict(checkpoint["model_state_dict"])
            model.eval()
            logger.info(f"Successfully loaded model from {model_path}")
            return model
        except Exception as exc:  # pragma: no cover - logging path
            logger.error(f"Failed to load model from {model_path}: {exc}")
            raise

    def predict(self, test_data_input) -> Iterator[QuantileForecast]:
        """Generate forecasts for the test data."""

        if hasattr(test_data_input, "__iter__") and not isinstance(test_data_input, list):
            test_data_input = list(test_data_input)
        logger.debug(f"Processing {len(test_data_input)} time series")

        # Group series by their effective length (after optional truncation),
        # then process each uniform-length group in sub-batches up to batch_size.
        def _effective_length(entry) -> int:
            target = entry["target"]
            if target.ndim == 1:
                seq_len = len(target)
            else:
                # target shape is [num_channels, seq_len]
                seq_len = target.shape[1]
            if self.max_context_length is not None:
                seq_len = min(seq_len, self.max_context_length)
            return seq_len

        length_to_items: dict[int, list[tuple[int, object]]] = {}
        for idx, entry in enumerate(test_data_input):
            seq_len = _effective_length(entry)
            length_to_items.setdefault(seq_len, []).append((idx, entry))

        total = len(test_data_input)
        ordered_results: list[QuantileForecast | None] = [None] * total

        for _, items in length_to_items.items():
            for i in range(0, len(items), self.batch_size):
                chunk = items[i : i + self.batch_size]
                entries = [entry for (_orig_idx, entry) in chunk]
                batch_forecasts = self._predict_batch(entries)
                for forecast_idx, (orig_idx, _entry) in enumerate(chunk):
                    ordered_results[orig_idx] = batch_forecasts[forecast_idx]

        return ordered_results  # type: ignore[return-value]

    def _predict_batch(self, test_data_batch: list) -> list[QuantileForecast]:
        """Generate predictions for a batch of time series."""

        logger.debug(f"Processing batch of size: {len(test_data_batch)}")

        try:
            batch_container = self._convert_to_batch_container(test_data_batch)

            if isinstance(device, torch.device):
                device_type = device.type
            else:
                device_type = "cuda" if "cuda" in str(device).lower() else "cpu"
            enable_autocast = device_type == "cuda"

            with torch.autocast(
                device_type=device_type,
                dtype=torch.bfloat16,
                enabled=enable_autocast,
            ):
                with torch.no_grad():
                    model_output = self.model(batch_container, drop_enc_allow=False)

            forecasts = self._convert_to_forecasts(model_output, test_data_batch, batch_container)

            logger.debug(f"Generated {len(forecasts)} forecasts")
            return forecasts
        except Exception as exc:  # pragma: no cover - logging path
            logger.error(f"Error in batch prediction: {exc}")
            raise

    def _convert_to_batch_container(self, test_data_batch: list) -> BatchTimeSeriesContainer:
        """Convert gluonts test data to BatchTimeSeriesContainer."""

        batch_size = len(test_data_batch)
        history_values_list = []
        start_dates = []
        frequencies = []

        for entry in test_data_batch:
            target = entry["target"]

            if target.ndim == 1:
                target = target.reshape(-1, 1)
            else:
                target = target.T

            if self.max_context_length is not None and len(target) > self.max_context_length:
                target = target[-self.max_context_length :]

            history_values_list.append(target)
            start_dates.append(entry["start"].to_timestamp().to_datetime64())
            frequencies.append(parse_frequency(entry["freq"]))

        history_values_np = np.stack(history_values_list, axis=0)
        num_channels = history_values_np.shape[2]

        history_values = torch.tensor(history_values_np, dtype=torch.float32, device=device)

        future_values = torch.zeros(
            (batch_size, self.ds_prediction_length, num_channels),
            dtype=torch.float32,
            device=device,
        )

        return BatchTimeSeriesContainer(
            history_values=history_values,
            future_values=future_values,
            start=start_dates,
            frequency=frequencies,
        )

    def _convert_to_forecasts(
        self,
        model_output: dict,
        test_data_batch: list,
        batch_container: BatchTimeSeriesContainer,
    ) -> list[QuantileForecast]:
        """Convert model predictions to QuantileForecast objects."""

        predictions = model_output["result"]
        scale_statistics = model_output["scale_statistics"]

        if predictions.ndim == 4:
            predictions_unscaled = self.scaler.inverse_scale(predictions, scale_statistics)
            is_quantile = True
            quantile_levels = self.model.quantiles
        else:
            predictions_unscaled = self.scaler.inverse_scale(predictions, scale_statistics)
            is_quantile = False
            quantile_levels = [0.5]

        forecasts: list[QuantileForecast] = []
        for idx, entry in enumerate(test_data_batch):
            history_length = int(batch_container.history_values.shape[1])
            start_date = entry["start"]
            forecast_start = start_date + history_length

            if is_quantile:
                pred_array = predictions_unscaled[idx].cpu().numpy()

                if pred_array.shape[1] == 1:
                    pred_array = pred_array.squeeze(1)
                    forecast_arrays = pred_array.T
                else:
                    forecast_arrays = pred_array.transpose(2, 0, 1)

                forecast = QuantileForecast(
                    forecast_arrays=forecast_arrays,
                    forecast_keys=[str(q) for q in quantile_levels],
                    start_date=forecast_start,
                )
            else:
                pred_array = predictions_unscaled[idx].cpu().numpy()

                if pred_array.shape[1] == 1:
                    pred_array = pred_array.squeeze(1)
                    forecast_arrays = pred_array.reshape(1, -1)
                else:
                    forecast_arrays = pred_array.reshape(1, *pred_array.shape)

                forecast = QuantileForecast(
                    forecast_arrays=forecast_arrays,
                    forecast_keys=["0.5"],
                    start_date=forecast_start,
                )

            forecasts.append(forecast)

        return forecasts


__all__ = ["TimeSeriesPredictor"]