"""SARIMA (Seasonal ARIMA) forecaster wrapper."""

from typing import Optional, Union, Tuple, List

import numpy as np
import pandas as pd
from scipy import stats

from ad_inventory_forecast.core.base import BaseForecaster

try:
    from statsmodels.tsa.statespace.sarimax import SARIMAX
    from statsmodels.tsa.stattools import adfuller, kpss
    HAS_STATSMODELS = True
except ImportError:
    HAS_STATSMODELS = False

try:
    import pmdarima as pm
    HAS_PMDARIMA = True
except ImportError:
    HAS_PMDARIMA = False


class SARIMAForecaster(BaseForecaster):
    """
    SARIMA (Seasonal ARIMA) forecaster.

    Wraps statsmodels SARIMAX with optional automatic order selection
    via pmdarima's auto_arima.

    SARIMA(p,d,q)(P,D,Q)[m] models combine:
    - AR(p): Autoregressive terms
    - I(d): Integration (differencing) order
    - MA(q): Moving average terms
    - Seasonal AR(P), I(D), MA(Q) with period m

    Parameters
    ----------
    order : Tuple[int, int, int], optional
        ARIMA order (p, d, q). If None and auto=True, automatically selected.
    seasonal_order : Tuple[int, int, int, int], optional
        Seasonal ARIMA order (P, D, Q, m). If None and auto=True, automatically selected.
    auto : bool, default True
        Whether to automatically select orders using pmdarima.
    max_p : int, default 3
        Maximum p for auto selection.
    max_d : int, default 2
        Maximum d for auto selection.
    max_q : int, default 3
        Maximum q for auto selection.
    max_P : int, default 2
        Maximum seasonal P for auto selection.
    max_D : int, default 1
        Maximum seasonal D for auto selection.
    max_Q : int, default 2
        Maximum seasonal Q for auto selection.
    seasonal_periods : int, optional
        Seasonal period m. Auto-detected if not specified.
    information_criterion : str, default 'aic'
        Criterion for model selection: 'aic', 'bic', 'aicc'.
    stepwise : bool, default True
        Whether to use stepwise search (faster) or exhaustive search.
    enforce_stationarity : bool, default True
        Whether to enforce stationarity in AR parameters.
    enforce_invertibility : bool, default True
        Whether to enforce invertibility in MA parameters.

    Attributes
    ----------
    is_fitted_ : bool
        Whether the model has been fitted.
    model_ : SARIMAX
        Fitted statsmodels model.
    result_ : SARIMAXResults
        Fitted model results.
    order_ : Tuple[int, int, int]
        Selected/specified ARIMA order.
    seasonal_order_ : Tuple[int, int, int, int]
        Selected/specified seasonal order.
    aic_ : float
        Akaike Information Criterion.
    bic_ : float
        Bayesian Information Criterion.

    Examples
    --------
    >>> import pandas as pd
    >>> import numpy as np
    >>> dates = pd.date_range('2024-01-01', periods=365, freq='D')
    >>> y = pd.Series(100 + np.sin(np.arange(365) * 2 * np.pi / 7) * 10, index=dates)
    >>> forecaster = SARIMAForecaster(auto=True, seasonal_periods=7)
    >>> forecaster.fit(y)
    >>> forecast = forecaster.predict(horizon=30)
    >>> print(f"Order: {forecaster.order_}, Seasonal: {forecaster.seasonal_order_}")

    Notes
    -----
    - Automatic order selection requires pmdarima to be installed
    - For large datasets, use stepwise=True for faster fitting
    - The model assumes data is non-negative (common for ad impressions)
    """

    def __init__(
        self,
        order: Optional[Tuple[int, int, int]] = None,
        seasonal_order: Optional[Tuple[int, int, int, int]] = None,
        auto: bool = True,
        max_p: int = 3,
        max_d: int = 2,
        max_q: int = 3,
        max_P: int = 2,
        max_D: int = 1,
        max_Q: int = 2,
        seasonal_periods: Optional[int] = None,
        information_criterion: str = "aic",
        stepwise: bool = True,
        enforce_stationarity: bool = True,
        enforce_invertibility: bool = True,
    ):
        """Initialize the SARIMA forecaster."""
        super().__init__()

        if not HAS_STATSMODELS:
            raise ImportError(
                "statsmodels is required for SARIMAForecaster. "
                "Install with: pip install statsmodels"
            )

        self.order = order
        self.seasonal_order = seasonal_order
        self.auto = auto
        self.max_p = max_p
        self.max_d = max_d
        self.max_q = max_q
        self.max_P = max_P
        self.max_D = max_D
        self.max_Q = max_Q
        self.seasonal_periods = seasonal_periods
        self.information_criterion = information_criterion
        self.stepwise = stepwise
        self.enforce_stationarity = enforce_stationarity
        self.enforce_invertibility = enforce_invertibility

        # Fitted attributes
        self.model_ = None
        self.result_ = None
        self.order_ = None
        self.seasonal_order_ = None
        self.aic_ = None
        self.bic_ = None
        self._last_date = None

    def fit(
        self,
        y: Union[pd.Series, np.ndarray],
        X: Optional[Union[pd.DataFrame, np.ndarray]] = None,
    ) -> "SARIMAForecaster":
        """
        Fit the SARIMA model to training data.

        Parameters
        ----------
        y : pd.Series or np.ndarray
            Target time series.
        X : pd.DataFrame or np.ndarray, optional
            Exogenous features (passed to SARIMAX).

        Returns
        -------
        self : SARIMAForecaster
            Fitted forecaster instance.
        """
        # Validate input
        y = self._validate_y(y)
        self.training_series_ = y.copy()
        self.n_obs_ = len(y)
        self.freq_ = self._infer_frequency(y)
        self._last_date = y.index[-1]

        # Determine seasonal periods if not specified
        if self.seasonal_periods is None:
            self.seasonal_periods = self._get_seasonal_period()

        # Determine model orders
        if self.auto and self.order is None:
            self._auto_select_orders(y)
        else:
            self.order_ = self.order or (1, 1, 1)
            if self.seasonal_order is not None:
                self.seasonal_order_ = self.seasonal_order
            else:
                # Default seasonal order
                self.seasonal_order_ = (1, 1, 1, self.seasonal_periods)

        # Fit SARIMAX model
        try:
            self.model_ = SARIMAX(
                y,
                order=self.order_,
                seasonal_order=self.seasonal_order_,
                enforce_stationarity=self.enforce_stationarity,
                enforce_invertibility=self.enforce_invertibility,
            )

            self.result_ = self.model_.fit(disp=False)
            self.aic_ = self.result_.aic
            self.bic_ = self.result_.bic

        except Exception as e:
            # Fall back to simpler model
            self.order_ = (1, 1, 0)
            self.seasonal_order_ = (0, 0, 0, self.seasonal_periods)

            self.model_ = SARIMAX(
                y,
                order=self.order_,
                seasonal_order=self.seasonal_order_,
            )
            self.result_ = self.model_.fit(disp=False)

        self.is_fitted_ = True
        return self

    def predict(
        self,
        horizon: int,
        return_conf_int: bool = False,
        alpha: float = 0.05,
        return_df: bool = True,
    ) -> Union[pd.Series, pd.DataFrame, Tuple[pd.Series, pd.DataFrame]]:
        """
        Generate forecasts for the specified horizon.

        Parameters
        ----------
        horizon : int
            Number of periods to forecast.
        return_conf_int : bool, default False
            Whether to return confidence intervals (legacy mode).
        alpha : float, default 0.05
            Significance level for confidence intervals.
        return_df : bool, default True
            If True (default), returns DataFrame with 'predicted_impressions',
            'lower_bound', 'upper_bound' columns (SDK-compliant format).

        Returns
        -------
        forecast_df : pd.DataFrame
            If return_df=True (default): DataFrame with prediction columns.
        forecasts : pd.Series
            If return_df=False: Point forecasts with DatetimeIndex.
        conf_int : pd.DataFrame, optional
            Confidence intervals with 'lower' and 'upper' columns.
        """
        self._check_is_fitted()

        # Generate forecast index
        forecast_index = self._generate_forecast_index(self._last_date, horizon)

        # Get forecast with confidence intervals
        forecast_result = self.result_.get_forecast(steps=horizon)
        forecast_values = forecast_result.predicted_mean

        # Ensure non-negative
        forecast_values = np.maximum(forecast_values.values, 0)

        forecast = pd.Series(
            forecast_values,
            index=forecast_index,
            name="forecast",
        )

        # Get confidence intervals
        conf_int_raw = forecast_result.conf_int(alpha=alpha)
        lower = np.maximum(conf_int_raw.iloc[:, 0].values, 0)
        upper = conf_int_raw.iloc[:, 1].values

        conf_int = pd.DataFrame(
            {"lower": lower, "upper": upper},
            index=forecast_index,
        )

        # Return SDK-compliant DataFrame by default
        if return_df and not return_conf_int:
            return pd.DataFrame({
                "predicted_impressions": forecast.values,
                "lower_bound": conf_int["lower"].values,
                "upper_bound": conf_int["upper"].values,
            }, index=forecast_index)

        if return_conf_int:
            return forecast, conf_int

        return forecast

    def _auto_select_orders(self, y: pd.Series) -> None:
        """Automatically select SARIMA orders using pmdarima."""
        if HAS_PMDARIMA:
            # Use pmdarima's auto_arima
            auto_model = pm.auto_arima(
                y.values,
                start_p=0,
                max_p=self.max_p,
                start_q=0,
                max_q=self.max_q,
                d=None,
                max_d=self.max_d,
                start_P=0,
                max_P=self.max_P,
                start_Q=0,
                max_Q=self.max_Q,
                D=None,
                max_D=self.max_D,
                m=self.seasonal_periods,
                seasonal=True if self.seasonal_periods > 1 else False,
                information_criterion=self.information_criterion,
                stepwise=self.stepwise,
                suppress_warnings=True,
                error_action="ignore",
                trace=False,
            )

            self.order_ = auto_model.order
            self.seasonal_order_ = auto_model.seasonal_order

        else:
            # Manual selection using ADF test and grid search
            self._manual_order_selection(y)

    def _manual_order_selection(self, y: pd.Series) -> None:
        """Manual order selection when pmdarima is not available."""
        # Determine differencing order using ADF test
        d = self._determine_differencing(y)

        # Simple grid search for p and q
        best_aic = np.inf
        best_order = (1, d, 1)

        for p in range(self.max_p + 1):
            for q in range(self.max_q + 1):
                if p == 0 and q == 0:
                    continue

                try:
                    model = SARIMAX(
                        y,
                        order=(p, d, q),
                        seasonal_order=(0, 0, 0, self.seasonal_periods),
                    )
                    result = model.fit(disp=False)

                    if result.aic < best_aic:
                        best_aic = result.aic
                        best_order = (p, d, q)

                except Exception:
                    continue

        self.order_ = best_order

        # Determine seasonal orders
        D = 1 if self.seasonal_periods > 1 else 0
        self.seasonal_order_ = (1, D, 1, self.seasonal_periods)

    def _determine_differencing(self, y: pd.Series) -> int:
        """Determine differencing order using stationarity tests."""
        # ADF test
        try:
            adf_result = adfuller(y.values, autolag="AIC")
            adf_pvalue = adf_result[1]
        except Exception:
            adf_pvalue = 1.0

        # KPSS test
        try:
            kpss_result = kpss(y.values, regression="c")
            kpss_pvalue = kpss_result[1]
        except Exception:
            kpss_pvalue = 0.0

        # Decision logic
        if adf_pvalue < 0.05 and kpss_pvalue > 0.05:
            # Stationary
            return 0
        elif adf_pvalue >= 0.05 and kpss_pvalue <= 0.05:
            # Non-stationary, difference once
            return 1
        else:
            # Ambiguous, try differencing
            y_diff = y.diff().dropna()
            try:
                adf_diff = adfuller(y_diff.values, autolag="AIC")[1]
                if adf_diff < 0.05:
                    return 1
                else:
                    return 2
            except Exception:
                return 1

    def get_diagnostics(self) -> dict:
        """
        Get diagnostic information about the fitted model.

        Returns
        -------
        diagnostics : dict
            Dictionary with model diagnostics.
        """
        self._check_is_fitted()

        return {
            "order": self.order_,
            "seasonal_order": self.seasonal_order_,
            "aic": self.aic_,
            "bic": self.bic_,
            "log_likelihood": self.result_.llf,
            "n_params": len(self.result_.params),
            "residual_mean": self.result_.resid.mean(),
            "residual_std": self.result_.resid.std(),
        }

    def get_residuals(self) -> pd.Series:
        """
        Get model residuals.

        Returns
        -------
        residuals : pd.Series
            Model residuals.
        """
        self._check_is_fitted()
        return pd.Series(
            self.result_.resid,
            index=self.training_series_.index,
            name="residuals",
        )

    def plot_diagnostics(self):
        """
        Plot model diagnostics.

        Returns
        -------
        fig : matplotlib.figure.Figure
            Diagnostic plots.
        """
        self._check_is_fitted()

        try:
            import matplotlib.pyplot as plt
        except ImportError:
            raise ImportError("matplotlib is required for plotting")

        fig = self.result_.plot_diagnostics(figsize=(12, 10))
        return fig

    def summary(self) -> str:
        """
        Get model summary.

        Returns
        -------
        summary : str
            Model summary string.
        """
        self._check_is_fitted()
        return str(self.result_.summary())


def auto_arima(
    y: pd.Series,
    seasonal_periods: Optional[int] = None,
    information_criterion: str = "aic",
    stepwise: bool = True,
) -> SARIMAForecaster:
    """
    Automatically select and fit the best SARIMA model.

    Parameters
    ----------
    y : pd.Series
        Time series to fit.
    seasonal_periods : int, optional
        Seasonal period. Auto-detected if not specified.
    information_criterion : str, default 'aic'
        Criterion for model selection.
    stepwise : bool, default True
        Whether to use stepwise search.

    Returns
    -------
    model : SARIMAForecaster
        Best fitted SARIMA model.

    Examples
    --------
    >>> model = auto_arima(y, seasonal_periods=7)
    >>> forecast = model.predict(horizon=30)
    """
    model = SARIMAForecaster(
        auto=True,
        seasonal_periods=seasonal_periods,
        information_criterion=information_criterion,
        stepwise=stepwise,
    )
    model.fit(y)
    return model
