"""Prophet forecaster wrapper."""

from typing import Optional, Union, Tuple, List, Dict, Any

import numpy as np
import pandas as pd

from ad_inventory_forecast.core.base import BaseForecaster

try:
    from prophet import Prophet
    HAS_PROPHET = True
except ImportError:
    HAS_PROPHET = False


class ProphetForecaster(BaseForecaster):
    """
    Facebook Prophet forecaster wrapper.

    Wraps the Prophet library with a scikit-learn compatible interface.
    Prophet is particularly good for:
    - Data with strong seasonal effects
    - Multiple seasonality (daily, weekly, yearly)
    - Holiday effects
    - Handling missing data and outliers
    - Automatic changepoint detection

    Parameters
    ----------
    yearly_seasonality : bool or int, default 'auto'
        Include yearly seasonality. If int, number of Fourier terms.
    weekly_seasonality : bool or int, default 'auto'
        Include weekly seasonality. If int, number of Fourier terms.
    daily_seasonality : bool or int, default False
        Include daily seasonality. If int, number of Fourier terms.
    seasonality_mode : {'additive', 'multiplicative'}, default 'multiplicative'
        Mode for seasonality component.
    holidays : pd.DataFrame, optional
        DataFrame with 'holiday' and 'ds' columns for holiday effects.
    changepoint_prior_scale : float, default 0.05
        Flexibility of automatic changepoint selection.
        Larger values = more changepoints = more flexible trend.
    seasonality_prior_scale : float, default 10.0
        Strength of seasonality model.
    holidays_prior_scale : float, default 10.0
        Strength of holiday effects.
    n_changepoints : int, default 25
        Number of potential changepoints to consider.
    changepoint_range : float, default 0.8
        Proportion of history for fitting changepoints.
    growth : {'linear', 'logistic', 'flat'}, default 'linear'
        Growth trend type.
    cap : float, optional
        Maximum capacity for logistic growth.
    floor : float, optional
        Minimum capacity for logistic growth.
    interval_width : float, default 0.95
        Width of prediction intervals.
    uncertainty_samples : int, default 1000
        Number of samples for uncertainty estimation.

    Attributes
    ----------
    is_fitted_ : bool
        Whether the model has been fitted.
    model_ : Prophet
        Fitted Prophet model.
    train_df_ : pd.DataFrame
        Training data in Prophet format.

    Examples
    --------
    >>> import pandas as pd
    >>> import numpy as np
    >>> dates = pd.date_range('2023-01-01', periods=365, freq='D')
    >>> y = pd.Series(100 + np.sin(np.arange(365) * 2 * np.pi / 7) * 10, index=dates)
    >>> forecaster = ProphetForecaster(weekly_seasonality=True)
    >>> forecaster.fit(y)
    >>> forecast = forecaster.predict(horizon=30)

    Notes
    -----
    Prophet requires the input to be a DataFrame with 'ds' (datetime) and
    'y' (value) columns. This wrapper handles the conversion automatically.
    """

    def __init__(
        self,
        yearly_seasonality: Union[bool, int] = "auto",
        weekly_seasonality: Union[bool, int] = "auto",
        daily_seasonality: Union[bool, int] = False,
        seasonality_mode: str = "multiplicative",
        holidays: Optional[pd.DataFrame] = None,
        changepoint_prior_scale: float = 0.05,
        seasonality_prior_scale: float = 10.0,
        holidays_prior_scale: float = 10.0,
        n_changepoints: int = 25,
        changepoint_range: float = 0.8,
        growth: str = "linear",
        cap: Optional[float] = None,
        floor: Optional[float] = None,
        interval_width: float = 0.95,
        uncertainty_samples: int = 1000,
    ):
        """Initialize the Prophet forecaster."""
        super().__init__()

        if not HAS_PROPHET:
            raise ImportError(
                "prophet is required for ProphetForecaster. "
                "Install with: pip install prophet"
            )

        self.yearly_seasonality = yearly_seasonality
        self.weekly_seasonality = weekly_seasonality
        self.daily_seasonality = daily_seasonality
        self.seasonality_mode = seasonality_mode
        self.holidays = holidays
        self.changepoint_prior_scale = changepoint_prior_scale
        self.seasonality_prior_scale = seasonality_prior_scale
        self.holidays_prior_scale = holidays_prior_scale
        self.n_changepoints = n_changepoints
        self.changepoint_range = changepoint_range
        self.growth = growth
        self.cap = cap
        self.floor = floor
        self.interval_width = interval_width
        self.uncertainty_samples = uncertainty_samples

        # Fitted attributes
        self.model_ = None
        self.train_df_ = None
        self._last_date = None

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

        Parameters
        ----------
        y : pd.Series or np.ndarray
            Target time series.
        X : pd.DataFrame or np.ndarray, optional
            Exogenous features (added as regressors).

        Returns
        -------
        self : ProphetForecaster
            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]

        # Convert to Prophet format
        self.train_df_ = pd.DataFrame({
            "ds": y.index,
            "y": y.values,
        })

        # Add cap/floor for logistic growth
        if self.growth == "logistic":
            if self.cap is not None:
                self.train_df_["cap"] = self.cap
            else:
                # Auto-set cap to 1.5x max value
                self.train_df_["cap"] = y.max() * 1.5

            if self.floor is not None:
                self.train_df_["floor"] = self.floor
            else:
                self.train_df_["floor"] = 0

        # Create Prophet model
        self.model_ = Prophet(
            yearly_seasonality=self.yearly_seasonality,
            weekly_seasonality=self.weekly_seasonality,
            daily_seasonality=self.daily_seasonality,
            seasonality_mode=self.seasonality_mode,
            holidays=self.holidays,
            changepoint_prior_scale=self.changepoint_prior_scale,
            seasonality_prior_scale=self.seasonality_prior_scale,
            holidays_prior_scale=self.holidays_prior_scale,
            n_changepoints=self.n_changepoints,
            changepoint_range=self.changepoint_range,
            growth=self.growth,
            interval_width=self.interval_width,
            uncertainty_samples=self.uncertainty_samples,
        )

        # Add exogenous regressors if provided
        if X is not None:
            if isinstance(X, pd.DataFrame):
                for col in X.columns:
                    self.model_.add_regressor(col)
                    self.train_df_[col] = X[col].values

        # Fit model (suppress output)
        import logging
        logging.getLogger("prophet").setLevel(logging.WARNING)
        logging.getLogger("cmdstanpy").setLevel(logging.WARNING)

        self.model_.fit(self.train_df_)

        self.is_fitted_ = True
        return self

    def predict(
        self,
        horizon: int,
        return_conf_int: bool = False,
        alpha: float = 0.05,
        return_df: bool = True,
        X_future: Optional[pd.DataFrame] = None,
    ) -> 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.
            Note: Prophet uses interval_width from initialization.
        return_df : bool, default True
            If True (default), returns DataFrame with 'predicted_impressions',
            'lower_bound', 'upper_bound' columns (SDK-compliant format).
        X_future : pd.DataFrame, optional
            Future values of exogenous regressors.

        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)

        # Create future DataFrame
        future_df = pd.DataFrame({"ds": forecast_index})

        # Add cap/floor for logistic growth
        if self.growth == "logistic":
            future_df["cap"] = self.cap or self.train_df_["cap"].iloc[0]
            future_df["floor"] = self.floor or 0

        # Add future regressors if provided
        if X_future is not None:
            for col in X_future.columns:
                future_df[col] = X_future[col].values

        # Generate forecast
        forecast = self.model_.predict(future_df)

        # Extract point forecasts
        forecast_values = forecast["yhat"].values

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

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

        # Get confidence intervals
        lower = np.maximum(forecast["yhat_lower"].values, 0)
        upper = forecast["yhat_upper"].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": lower,
                "upper_bound": upper,
            }, index=forecast_index)

        if return_conf_int:
            return result, conf_int

        return result

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

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

        # Get forecast for training period
        forecast = self.model_.predict(self.train_df_)

        components = pd.DataFrame(index=self.training_series_.index)
        components["trend"] = forecast["trend"].values

        # Add seasonality components
        for col in forecast.columns:
            if col.endswith("_weekly") or col.endswith("_yearly") or col.endswith("_daily"):
                components[col] = forecast[col].values

        components["fitted"] = forecast["yhat"].values
        components["residual"] = self.training_series_.values - forecast["yhat"].values

        return components

    def get_changepoints(self) -> pd.DatetimeIndex:
        """
        Get detected changepoints.

        Returns
        -------
        changepoints : pd.DatetimeIndex
            Dates of detected trend changepoints.
        """
        self._check_is_fitted()
        return pd.DatetimeIndex(self.model_.changepoints)

    def add_seasonality(
        self,
        name: str,
        period: float,
        fourier_order: int,
        prior_scale: float = 10.0,
        mode: Optional[str] = None,
    ) -> "ProphetForecaster":
        """
        Add custom seasonality before fitting.

        Parameters
        ----------
        name : str
            Name of the seasonality component.
        period : float
            Period of seasonality in days.
        fourier_order : int
            Number of Fourier terms.
        prior_scale : float, default 10.0
            Strength of seasonality.
        mode : str, optional
            'additive' or 'multiplicative'. Uses model default if None.

        Returns
        -------
        self : ProphetForecaster
            Self for method chaining.
        """
        if self.is_fitted_:
            raise ValueError("Cannot add seasonality after fitting")

        # Store for later addition
        if not hasattr(self, "_custom_seasonalities"):
            self._custom_seasonalities = []

        self._custom_seasonalities.append({
            "name": name,
            "period": period,
            "fourier_order": fourier_order,
            "prior_scale": prior_scale,
            "mode": mode or self.seasonality_mode,
        })

        return self

    def plot_components(self):
        """
        Plot model components.

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

        forecast = self.model_.predict(self.train_df_)
        fig = self.model_.plot_components(forecast)
        return fig

    def plot_forecast(
        self,
        forecast: pd.Series,
        conf_int: Optional[pd.DataFrame] = None,
        history_periods: int = 90,
    ):
        """
        Plot forecast with historical data.

        Parameters
        ----------
        forecast : pd.Series
            Forecast values.
        conf_int : pd.DataFrame, optional
            Confidence intervals.
        history_periods : int, default 90
            Number of historical periods to show.

        Returns
        -------
        fig : matplotlib.figure.Figure
            Forecast plot.
        """
        self._check_is_fitted()

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

        fig, ax = plt.subplots(figsize=(12, 6))

        # Plot historical data
        history = self.training_series_.iloc[-history_periods:]
        ax.plot(history.index, history.values, "k.", label="Historical", alpha=0.5)

        # Plot forecast
        ax.plot(forecast.index, forecast.values, "b-", label="Forecast", linewidth=2)

        # Plot confidence intervals
        if conf_int is not None:
            ax.fill_between(
                forecast.index,
                conf_int["lower"],
                conf_int["upper"],
                color="blue",
                alpha=0.2,
                label="95% CI",
            )

        ax.set_title("Prophet Forecast")
        ax.set_xlabel("Date")
        ax.set_ylabel("Value")
        ax.legend()

        plt.tight_layout()
        return fig


def create_holiday_df(
    holidays: Dict[str, List[str]],
    lower_window: int = 0,
    upper_window: int = 1,
) -> pd.DataFrame:
    """
    Create a holiday DataFrame for Prophet.

    Parameters
    ----------
    holidays : Dict[str, List[str]]
        Dictionary mapping holiday names to lists of date strings.
        Example: {'super_bowl': ['2023-02-12', '2024-02-11']}
    lower_window : int, default 0
        Days before holiday to include.
    upper_window : int, default 1
        Days after holiday to include.

    Returns
    -------
    holiday_df : pd.DataFrame
        DataFrame formatted for Prophet's holidays parameter.

    Examples
    --------
    >>> holidays = {
    ...     'super_bowl': ['2023-02-12', '2024-02-11'],
    ...     'black_friday': ['2023-11-24', '2024-11-29'],
    ... }
    >>> holiday_df = create_holiday_df(holidays)
    """
    records = []

    for name, dates in holidays.items():
        for date_str in dates:
            records.append({
                "holiday": name,
                "ds": pd.Timestamp(date_str),
                "lower_window": lower_window,
                "upper_window": upper_window,
            })

    return pd.DataFrame(records)
