"""ETS (Exponential Smoothing) forecaster wrapper."""

from typing import Optional, Union, Tuple, Literal

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

from ad_inventory_forecast.core.base import BaseForecaster
from ad_inventory_forecast.core.diagnostics import (
    analyze_volatility,
    VolatilityAnalysis,
)

try:
    from statsmodels.tsa.holtwinters import ExponentialSmoothing
    HAS_STATSMODELS = True
except ImportError:
    HAS_STATSMODELS = False


class ETSForecaster(BaseForecaster):
    """
    Exponential Smoothing (ETS) forecaster.

    Wraps statsmodels ExponentialSmoothing with automatic model selection
    and scikit-learn compatible interface.

    ETS models decompose a time series into Error, Trend, and Seasonal
    components. Common configurations:
    - ETS(A,N,N): Simple exponential smoothing (no trend, no seasonality)
    - ETS(A,A,N): Holt's linear method (additive trend, no seasonality)
    - ETS(A,A,A): Holt-Winters additive (additive trend + seasonality)
    - ETS(A,A,M): Holt-Winters multiplicative seasonality
    - ETS(A,Ad,M): Damped trend with multiplicative seasonality

    Parameters
    ----------
    trend : {'add', 'mul', None, 'auto'}, default 'auto'
        Trend component type. 'auto' selects based on data.
    damped_trend : bool, default True
        Whether to use damped trend (recommended for forecasting).
    seasonal : {'add', 'mul', None, 'auto'}, default 'auto'
        Seasonal component type. 'auto' selects based on data.
    seasonal_periods : int, optional
        Number of periods in a seasonal cycle. Required if seasonal is not None.
        Auto-detected if not specified.
    initialization_method : str, default 'estimated'
        Method for initializing smoothing parameters.
        Options: 'estimated', 'heuristic', 'legacy-heuristic'.
    use_boxcox : bool or float, default False
        Whether to apply Box-Cox transformation.
        - False: No transformation
        - True: Estimate optimal lambda
        - float: Use specified lambda
    remove_bias : bool, default True
        Whether to remove bias from forecasts.

    Attributes
    ----------
    is_fitted_ : bool
        Whether the model has been fitted.
    model_ : ExponentialSmoothing
        Fitted statsmodels model.
    result_ : HoltWintersResults
        Fitted model results.
    aic_ : float
        Akaike Information Criterion of fitted model.
    bic_ : float
        Bayesian Information Criterion of fitted model.

    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 = ETSForecaster(seasonal='add', seasonal_periods=7)
    >>> forecaster.fit(y)
    >>> forecast = forecaster.predict(horizon=30)

    Notes
    -----
    - For multiplicative models, the series must be strictly positive
    - Damped trend is recommended for longer forecast horizons to prevent
      unrealistic growth
    - The 'auto' settings attempt to select the best model variant
    """

    def __init__(
        self,
        trend: Literal["add", "mul", None, "auto"] = "auto",
        damped_trend: bool = True,
        seasonal: Literal["add", "mul", None, "auto"] = "auto",
        seasonal_periods: Optional[int] = None,
        initialization_method: str = "estimated",
        use_boxcox: Union[bool, float] = False,
        remove_bias: bool = True,
    ):
        """Initialize the ETS forecaster."""
        super().__init__()

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

        self.trend = trend
        self.damped_trend = damped_trend
        self.seasonal = seasonal
        self.seasonal_periods = seasonal_periods
        self.initialization_method = initialization_method
        self.use_boxcox = use_boxcox
        self.remove_bias = remove_bias

        # Fitted attributes
        self.model_ = None
        self.result_ = None
        self.aic_ = None
        self.bic_ = None
        self._trend_type = None
        self._seasonal_type = None
        self._last_date = None
        self.volatility_: Optional[VolatilityAnalysis] = None

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

        Parameters
        ----------
        y : pd.Series or np.ndarray
            Target time series. Values must be positive for multiplicative
            models.
        X : pd.DataFrame or np.ndarray, optional
            Exogenous features (not used by ETS).

        Returns
        -------
        self : ETSForecaster
            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()

        # Select model components
        self._trend_type = self._select_trend(y)
        self._seasonal_type = self._select_seasonal(y)

        # Handle data requirements for multiplicative models
        y_fit = y.copy()
        if self._seasonal_type == "mul" or self._trend_type == "mul":
            if (y_fit <= 0).any():
                # Shift to make positive
                min_val = y_fit.min()
                if min_val <= 0:
                    y_fit = y_fit - min_val + 1

        # Create and fit model
        try:
            self.model_ = ExponentialSmoothing(
                y_fit,
                trend=self._trend_type,
                damped_trend=self.damped_trend if self._trend_type else False,
                seasonal=self._seasonal_type,
                seasonal_periods=self.seasonal_periods if self._seasonal_type else None,
                initialization_method=self.initialization_method,
                use_boxcox=self.use_boxcox,
            )

            self.result_ = self.model_.fit(
                optimized=True,
                remove_bias=self.remove_bias,
            )

            self.aic_ = self.result_.aic
            self.bic_ = self.result_.bic

        except Exception as e:
            # Fall back to simpler model
            self.model_ = ExponentialSmoothing(
                y_fit,
                trend="add",
                damped_trend=True,
                seasonal=None,
            )
            self.result_ = self.model_.fit(optimized=True)
            self._trend_type = "add"
            self._seasonal_type = None

        # Volatility analysis for interval widening
        self.volatility_ = analyze_volatility(y)

        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 point forecasts
        forecast_values = self.result_.forecast(horizon)

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

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

        # Calculate confidence intervals
        conf_int = self._calculate_confidence_intervals(
            forecast, horizon, alpha
        )

        # 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 _select_trend(self, y: pd.Series) -> Optional[str]:
        """Select trend component type."""
        if self.trend != "auto":
            return self.trend if self.trend else None

        # Check for trend using linear regression
        t = np.arange(len(y))
        slope, _, r_value, p_value, _ = stats.linregress(t, y.values)

        # Significant trend?
        if p_value < 0.05 and abs(r_value) > 0.3:
            # Check if multiplicative trend is better
            if (y > 0).all():
                log_slope, _, log_r, _, _ = stats.linregress(t, np.log(y.values))
                if abs(log_r) > abs(r_value):
                    return "mul"
            return "add"

        return None

    def _select_seasonal(self, y: pd.Series) -> Optional[str]:
        """Select seasonal component type based on data characteristics and length."""
        if self.seasonal != "auto":
            return self.seasonal if self.seasonal else None

        # Disable yearly seasonality for short history (< 365 days)
        # Only allow weekly or shorter seasonal periods
        if len(y) < 365 and self.seasonal_periods and self.seasonal_periods > 52:
            # Filter to weekly seasonality only
            self.seasonal_periods = min(self.seasonal_periods, 7)

        # Need enough data for seasonality (at least 2 full cycles)
        if len(y) < 2 * self.seasonal_periods:
            return None

        # Check for seasonality using autocorrelation
        from statsmodels.tsa.stattools import acf

        try:
            acf_values = acf(y.values, nlags=self.seasonal_periods + 1, fft=True)
            seasonal_acf = acf_values[self.seasonal_periods]

            # Significant seasonality?
            threshold = 2 / np.sqrt(len(y))
            if abs(seasonal_acf) > threshold:
                # Check if multiplicative is better
                if (y > 0).all():
                    # Compare coefficient of variation across seasons
                    n_seasons = len(y) // self.seasonal_periods
                    if n_seasons >= 2:
                        cv_values = []
                        for i in range(n_seasons):
                            season = y.iloc[
                                i * self.seasonal_periods : (i + 1) * self.seasonal_periods
                            ]
                            cv = season.std() / season.mean()
                            cv_values.append(cv)

                        # If CV is consistent, multiplicative might be better
                        cv_std = np.std(cv_values)
                        if cv_std < 0.1:
                            return "mul"

                return "add"

        except Exception:
            pass

        return None

    def _calculate_confidence_intervals(
        self,
        forecast: pd.Series,
        horizon: int,
        alpha: float,
    ) -> pd.DataFrame:
        """Calculate prediction intervals with volatility-based widening."""
        # Get residual standard error
        residuals = self.result_.resid
        sigma = np.std(residuals)

        # Calculate intervals with increasing uncertainty
        z = stats.norm.ppf(1 - alpha / 2)
        h = np.arange(1, horizon + 1)

        # Uncertainty grows with horizon
        uncertainty = sigma * np.sqrt(h)

        # Apply volatility multiplier for high-volatility series
        volatility_multiplier = 1.0
        if self.volatility_ is not None:
            volatility_multiplier = self.volatility_.interval_multiplier

        uncertainty = uncertainty * volatility_multiplier

        lower = forecast.values - z * uncertainty
        upper = forecast.values + z * uncertainty

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

        return pd.DataFrame(
            {"lower": lower, "upper": upper},
            index=forecast.index,
        )

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

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

        diagnostics = {
            "trend": self._trend_type,
            "damped_trend": self.damped_trend if self._trend_type else False,
            "seasonal": self._seasonal_type,
            "seasonal_periods": self.seasonal_periods,
            "aic": self.aic_,
            "bic": self.bic_,
            "smoothing_level": self.result_.params.get("smoothing_level"),
            "smoothing_trend": self.result_.params.get("smoothing_trend"),
            "smoothing_seasonal": self.result_.params.get("smoothing_seasonal"),
            "damping_trend": self.result_.params.get("damping_trend"),
            "sse": self.result_.sse,
        }

        # Add volatility diagnostics
        if self.volatility_ is not None:
            diagnostics["volatility_cv"] = self.volatility_.coefficient_of_variation
            diagnostics["volatility_level"] = self.volatility_.volatility_level
            diagnostics["interval_multiplier"] = self.volatility_.interval_multiplier

        return diagnostics

    def get_components(self) -> pd.DataFrame:
        """
        Get fitted components (level, trend, seasonal).

        Returns
        -------
        components : pd.DataFrame
            DataFrame with level, trend, and seasonal columns.
        """
        self._check_is_fitted()

        components = pd.DataFrame(index=self.training_series_.index)
        components["level"] = self.result_.level
        components["fitted"] = self.result_.fittedvalues
        components["residual"] = self.result_.resid

        if self._trend_type:
            # Trend is the slope/growth component
            components["trend"] = self.result_.trend

        if self._seasonal_type:
            components["seasonal"] = self.result_.season

        return components

    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, axes = plt.subplots(2, 2, figsize=(12, 10))

        # Fitted vs Actual
        axes[0, 0].plot(
            self.training_series_.index,
            self.training_series_.values,
            label="Actual",
            alpha=0.7,
        )
        axes[0, 0].plot(
            self.result_.fittedvalues.index,
            self.result_.fittedvalues.values,
            label="Fitted",
            alpha=0.7,
        )
        axes[0, 0].set_title("Fitted vs Actual")
        axes[0, 0].legend()

        # Residuals
        residuals = self.result_.resid
        axes[0, 1].plot(residuals.index, residuals.values)
        axes[0, 1].axhline(y=0, color="r", linestyle="--")
        axes[0, 1].set_title("Residuals")

        # Residual histogram
        axes[1, 0].hist(residuals.values, bins=30, edgecolor="black")
        axes[1, 0].set_title("Residual Distribution")

        # ACF of residuals
        from statsmodels.graphics.tsaplots import plot_acf

        plot_acf(residuals.values, ax=axes[1, 1], lags=min(40, len(residuals) // 2))
        axes[1, 1].set_title("ACF of Residuals")

        plt.tight_layout()
        return fig


def auto_ets(
    y: pd.Series,
    seasonal_periods: Optional[int] = None,
    information_criterion: str = "aic",
) -> ETSForecaster:
    """
    Automatically select the best ETS model.

    Tries multiple ETS configurations and returns the one with
    the best information criterion.

    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: 'aic' or 'bic'.

    Returns
    -------
    best_model : ETSForecaster
        Best fitted ETS model.

    Examples
    --------
    >>> model = auto_ets(y, seasonal_periods=7)
    >>> forecast = model.predict(horizon=30)
    """
    # Configurations to try
    configs = [
        {"trend": None, "seasonal": None},
        {"trend": "add", "seasonal": None},
        {"trend": "add", "damped_trend": True, "seasonal": None},
    ]

    # Add seasonal configurations if enough data
    if seasonal_periods is None:
        seasonal_periods = 7 if len(y) >= 14 else None

    if seasonal_periods and len(y) >= 2 * seasonal_periods:
        configs.extend([
            {"trend": None, "seasonal": "add", "seasonal_periods": seasonal_periods},
            {"trend": "add", "seasonal": "add", "seasonal_periods": seasonal_periods},
            {"trend": "add", "damped_trend": True, "seasonal": "add",
             "seasonal_periods": seasonal_periods},
        ])

        # Add multiplicative if data is positive
        if (y > 0).all():
            configs.extend([
                {"trend": None, "seasonal": "mul", "seasonal_periods": seasonal_periods},
                {"trend": "add", "seasonal": "mul", "seasonal_periods": seasonal_periods},
                {"trend": "add", "damped_trend": True, "seasonal": "mul",
                 "seasonal_periods": seasonal_periods},
            ])

    best_model = None
    best_ic = np.inf

    for config in configs:
        try:
            model = ETSForecaster(**config)
            model.fit(y)

            ic = model.aic_ if information_criterion == "aic" else model.bic_

            if ic < best_ic:
                best_ic = ic
                best_model = model

        except Exception:
            continue

    if best_model is None:
        # Fall back to simple exponential smoothing
        best_model = ETSForecaster(trend=None, seasonal=None)
        best_model.fit(y)

    return best_model
