"""Multi-segment forecasting pipeline."""

import warnings
from typing import List, Union

import numpy as np
import pandas as pd

from ad_inventory_forecast.models.auto import AutoForecaster


class ForecastPipeline:
    """Forecast a base metric across grouping dimensions from a raw tabular dataset.

    Parameters
    ----------
    df:
        Raw tabular DataFrame with one row per (date, segment) observation.
    time_col:
        Name of the column containing dates/timestamps.
    base_field:
        Name of the column containing the primary metric to forecast (e.g.
        ``"Ad Requests"``, ``"Impressions"``). Must be numeric and contain no
        infinite values. This parameter is required and must be specified
        explicitly — there is no default.
    strategy:
        AutoForecaster strategy — ``"fast"``, ``"balanced"``, or ``"thorough"``.
    min_obs:
        Minimum number of date observations required for a segment to be forecast.
        Segments below this threshold are skipped.
    verbose:
        If ``True``, emit warnings when segments are skipped.

    Raises
    ------
    ValueError
        If ``time_col`` is not found in ``df``, cannot be parsed as datetime,
        ``base_field`` is not found, is not numeric, or contains infinite values.
    """

    def __init__(
        self,
        df: pd.DataFrame,
        time_col: str,
        base_field: str,
        strategy: str = "fast",
        min_obs: int = 14,
        verbose: bool = False,
    ) -> None:
        if time_col not in df.columns:
            raise ValueError(
                f"time_col '{time_col}' not found in DataFrame. "
                f"Available columns: {list(df.columns)}"
            )

        # --- base_field validation ---
        if base_field not in df.columns:
            numeric_cols = df.select_dtypes(include="number").columns.tolist()
            raise ValueError(
                f"base_field '{base_field}' not found in DataFrame. "
                f"Available columns: {list(df.columns)}\n"
                f"Numeric columns (valid candidates): {numeric_cols}"
            )

        col = df[base_field]

        if not pd.api.types.is_numeric_dtype(col):
            raise ValueError(
                f"base_field '{base_field}' is not numeric "
                f"(dtype={col.dtype}). Choose a numeric column."
            )

        n_null = col.isna().sum()
        if n_null > 0:
            raise ValueError(
                f"base_field '{base_field}' contains {n_null} null value(s). "
                "Fill or drop nulls before constructing ForecastPipeline."
            )

        n_inf = np.isinf(col).sum()
        if n_inf > 0:
            raise ValueError(
                f"base_field '{base_field}' contains {n_inf} infinite value(s). "
                "Replace or drop them before constructing ForecastPipeline."
            )

        # --- time_col parsing ---
        try:
            df = df.copy()
            df[time_col] = pd.to_datetime(df[time_col])
        except Exception as exc:
            raise ValueError(
                f"Could not parse column '{time_col}' as datetime: {exc}"
            ) from exc

        self.df = df
        self.time_col = time_col
        self.base_field = base_field
        self.strategy = strategy
        self.min_obs = min_obs
        self.verbose = verbose

    def forecast(
        self,
        horizon: int,
        by: Union[None, str, List[str]] = None,
        alpha: float = 0.05,
    ) -> pd.DataFrame:
        """Produce forecasts for the given horizon.

        Parameters
        ----------
        horizon:
            Number of periods to forecast.
        by:
            Grouping column(s). ``None`` aggregates everything into a single series.
            A string or list of strings groups by those columns and returns one
            forecast block per segment.
        alpha:
            Significance level for prediction intervals (default 0.05 → 95 % CI).

        Returns
        -------
        pd.DataFrame
            When ``by=None``: ``(horizon, 3)`` DataFrame with
            ``predicted_impressions``, ``lower_bound``, ``upper_bound`` and a
            DatetimeIndex.

            When ``by`` is specified: flat DataFrame with segment columns prepended
            before the three forecast columns, indexed by date.
        """
        if by is None:
            return self._forecast_aggregate(horizon, alpha)

        by_cols = [by] if isinstance(by, str) else list(by)
        missing = [c for c in by_cols if c not in self.df.columns]
        if missing:
            raise ValueError(
                f"by column(s) {missing} not found in DataFrame. "
                f"Available columns: {list(self.df.columns)}"
            )

        return self._forecast_by(horizon, by_cols, alpha)

    # ------------------------------------------------------------------
    # Private helpers
    # ------------------------------------------------------------------

    def _fit_predict(self, series: pd.Series, horizon: int, alpha: float) -> pd.DataFrame:
        model = AutoForecaster(strategy=self.strategy)
        model.fit(series)
        return model.predict(horizon=horizon, alpha=alpha, return_df=True)

    def _aggregate_series(self, group: pd.DataFrame) -> pd.Series:
        series = group.groupby(self.time_col)[self.base_field].sum().sort_index()
        series.index = pd.DatetimeIndex(series.index)
        return series

    def _forecast_aggregate(self, horizon: int, alpha: float) -> pd.DataFrame:
        series = self._aggregate_series(self.df)
        return self._fit_predict(series, horizon, alpha)

    def _forecast_by(
        self, horizon: int, by_cols: List[str], alpha: float
    ) -> pd.DataFrame:
        frames = []

        for keys, group in self.df.groupby(by_cols):
            series = self._aggregate_series(group)

            if len(series) < self.min_obs:
                if self.verbose:
                    warnings.warn(
                        f"Skipping segment {keys}: only {len(series)} observations "
                        f"(min_obs={self.min_obs})"
                    )
                continue

            try:
                pred = self._fit_predict(series, horizon, alpha)
            except Exception as exc:
                if self.verbose:
                    warnings.warn(f"Skipping segment {keys}: {exc}")
                continue

            keys_tuple = keys if isinstance(keys, tuple) else (keys,)
            for col, val in zip(by_cols, keys_tuple):
                pred.insert(0, col, val)

            frames.append(pred)

        if not frames:
            raise ValueError(
                "No segments had enough observations to forecast. "
                f"Check min_obs={self.min_obs} or the data coverage."
            )

        return pd.concat(frames).sort_index()
