"""ML Ensemble forecaster using XGBoost/LightGBM."""

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

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

from ad_inventory_forecast.core.base import BaseForecaster

try:
    import xgboost as xgb
    HAS_XGBOOST = True
except ImportError:
    HAS_XGBOOST = False

try:
    import lightgbm as lgb
    HAS_LIGHTGBM = True
except ImportError:
    HAS_LIGHTGBM = False


class MLEnsembleForecaster(BaseForecaster):
    """
    Machine Learning ensemble forecaster using gradient boosting.

    Uses XGBoost or LightGBM with time series feature engineering
    for forecasting. Generates features from:
    - Temporal attributes (day_of_week, month, etc.)
    - Lag features (Y(t-1), Y(t-7), etc.)
    - Rolling statistics (mean, std over windows)
    - Trend features (time index)

    Parameters
    ----------
    model_type : {'xgboost', 'lightgbm'}, default 'xgboost'
        Which gradient boosting library to use.
    lag_features : List[int], default [1, 7, 14, 28]
        Lag periods to use as features.
    rolling_windows : List[int], default [7, 28]
        Window sizes for rolling statistics.
    n_estimators : int, default 100
        Number of boosting rounds.
    max_depth : int, default 6
        Maximum tree depth.
    learning_rate : float, default 0.1
        Learning rate.
    early_stopping_rounds : int, default 10
        Early stopping patience.
    validation_fraction : float, default 0.1
        Fraction of data for validation.
    random_state : int, optional
        Random seed for reproducibility.

    Attributes
    ----------
    is_fitted_ : bool
        Whether the model has been fitted.
    model_ : XGBRegressor or LGBMRegressor
        Fitted model.
    feature_names_ : List[str]
        Names of features used.
    feature_importance_ : pd.Series
        Feature importances.

    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 = MLEnsembleForecaster(model_type='xgboost')
    >>> forecaster.fit(y)
    >>> forecast = forecaster.predict(horizon=30)

    Notes
    -----
    ML models require sufficient history to create lag features.
    The minimum required history is max(lag_features) + max(rolling_windows).

    For forecasting multiple steps ahead, this uses recursive forecasting:
    predict one step, update features, repeat.
    """

    def __init__(
        self,
        model_type: Literal["xgboost", "lightgbm"] = "xgboost",
        lag_features: List[int] = None,
        rolling_windows: List[int] = None,
        n_estimators: int = 100,
        max_depth: int = 6,
        learning_rate: float = 0.1,
        early_stopping_rounds: int = 10,
        validation_fraction: float = 0.1,
        random_state: Optional[int] = None,
    ):
        """Initialize the ML ensemble forecaster."""
        super().__init__()

        self.model_type = model_type
        self.lag_features = lag_features or [1, 7, 14, 28]
        self.rolling_windows = rolling_windows or [7, 28]
        self.n_estimators = n_estimators
        self.max_depth = max_depth
        self.learning_rate = learning_rate
        self.early_stopping_rounds = early_stopping_rounds
        self.validation_fraction = validation_fraction
        self.random_state = random_state

        # Check dependencies
        if model_type == "xgboost" and not HAS_XGBOOST:
            raise ImportError(
                "xgboost is required for model_type='xgboost'. "
                "Install with: pip install xgboost"
            )
        if model_type == "lightgbm" and not HAS_LIGHTGBM:
            raise ImportError(
                "lightgbm is required for model_type='lightgbm'. "
                "Install with: pip install lightgbm"
            )

        # Fitted attributes
        self.model_ = None
        self.feature_names_ = None
        self.feature_importance_ = None
        self._last_date = None
        self._history_for_features = None

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

        Parameters
        ----------
        y : pd.Series or np.ndarray
            Target time series.
        X : pd.DataFrame or np.ndarray, optional
            Additional exogenous features.

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

        # Check minimum history
        min_history = max(self.lag_features) + max(self.rolling_windows)
        if len(y) < min_history + 10:
            raise ValueError(
                f"Need at least {min_history + 10} observations. "
                f"Got {len(y)}."
            )

        # Generate features
        X_features, y_target = self._generate_training_features(y)

        # Add external features if provided
        if X is not None:
            if isinstance(X, pd.DataFrame):
                # Align and merge
                X_aligned = X.loc[y_target.index]
                X_features = pd.concat([X_features, X_aligned], axis=1)
            else:
                # Assume aligned
                X_array = np.asarray(X)[min_history:]
                for i in range(X_array.shape[1] if X_array.ndim > 1 else 1):
                    col = X_array[:, i] if X_array.ndim > 1 else X_array
                    X_features[f"exog_{i}"] = col[:len(X_features)]

        self.feature_names_ = list(X_features.columns)

        # Split validation set
        n_val = int(len(X_features) * self.validation_fraction)
        if n_val < 10:
            n_val = 0

        X_train = X_features.iloc[:-n_val] if n_val > 0 else X_features
        y_train = y_target.iloc[:-n_val] if n_val > 0 else y_target
        X_val = X_features.iloc[-n_val:] if n_val > 0 else None
        y_val = y_target.iloc[-n_val:] if n_val > 0 else None

        # Create and fit model
        if self.model_type == "xgboost":
            self.model_ = xgb.XGBRegressor(
                n_estimators=self.n_estimators,
                max_depth=self.max_depth,
                learning_rate=self.learning_rate,
                random_state=self.random_state,
                verbosity=0,
            )

            if n_val > 0:
                self.model_.fit(
                    X_train,
                    y_train,
                    eval_set=[(X_val, y_val)],
                    verbose=False,
                )
            else:
                self.model_.fit(X_train, y_train)

        else:  # lightgbm
            self.model_ = lgb.LGBMRegressor(
                n_estimators=self.n_estimators,
                max_depth=self.max_depth,
                learning_rate=self.learning_rate,
                random_state=self.random_state,
                verbose=-1,
            )

            if n_val > 0:
                self.model_.fit(
                    X_train,
                    y_train,
                    eval_set=[(X_val, y_val)],
                )
            else:
                self.model_.fit(X_train, y_train)

        # Store feature importance
        self.feature_importance_ = pd.Series(
            self.model_.feature_importances_,
            index=self.feature_names_,
            name="importance",
        ).sort_values(ascending=False)

        # Store history for recursive forecasting
        self._history_for_features = y.copy()

        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 using recursive prediction.

        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)

        # Recursive forecasting
        history = self._history_for_features.copy()
        forecasts = []

        for i in range(horizon):
            # Generate features for next step
            X_next = self._generate_single_step_features(
                history, forecast_index[i]
            )

            # Predict
            pred = self.model_.predict(X_next)[0]
            pred = max(0, pred)  # Non-negative
            forecasts.append(pred)

            # Update history for next iteration
            new_row = pd.Series([pred], index=[forecast_index[i]])
            history = pd.concat([history, new_row])

        forecast_values = np.array(forecasts)

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

        # Calculate confidence intervals
        conf_int = self._calculate_confidence_intervals(
            result, 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 result, conf_int

        return result

    def _generate_training_features(
        self, y: pd.Series
    ) -> Tuple[pd.DataFrame, pd.Series]:
        """Generate feature matrix and aligned target for training."""
        features = pd.DataFrame(index=y.index)

        # Temporal features
        features["day_of_week"] = y.index.dayofweek
        features["day_of_month"] = y.index.day
        features["month"] = y.index.month
        features["week_of_year"] = y.index.isocalendar().week.values
        features["is_weekend"] = (y.index.dayofweek >= 5).astype(int)

        # Quarter and year
        features["quarter"] = y.index.quarter
        features["day_of_year"] = y.index.dayofyear

        # Lag features
        for lag in self.lag_features:
            features[f"lag_{lag}"] = y.shift(lag)

        # Rolling features
        for window in self.rolling_windows:
            features[f"rolling_mean_{window}"] = y.rolling(window=window).mean()
            features[f"rolling_std_{window}"] = y.rolling(window=window).std()
            features[f"rolling_min_{window}"] = y.rolling(window=window).min()
            features[f"rolling_max_{window}"] = y.rolling(window=window).max()

        # Trend feature (time index)
        features["trend"] = np.arange(len(y))

        # Drop rows with NaN (from lag/rolling features)
        min_rows = max(max(self.lag_features), max(self.rolling_windows))
        features = features.iloc[min_rows:]
        y_aligned = y.iloc[min_rows:]

        return features, y_aligned

    def _generate_single_step_features(
        self, history: pd.Series, target_date: pd.Timestamp
    ) -> pd.DataFrame:
        """Generate features for a single future step."""
        features = {}

        # Temporal features
        features["day_of_week"] = target_date.dayofweek
        features["day_of_month"] = target_date.day
        features["month"] = target_date.month
        features["week_of_year"] = target_date.isocalendar().week
        features["is_weekend"] = int(target_date.dayofweek >= 5)
        features["quarter"] = target_date.quarter
        features["day_of_year"] = target_date.dayofyear

        # Lag features
        for lag in self.lag_features:
            features[f"lag_{lag}"] = history.iloc[-lag] if lag <= len(history) else np.nan

        # Rolling features
        for window in self.rolling_windows:
            if len(history) >= window:
                recent = history.iloc[-window:]
                features[f"rolling_mean_{window}"] = recent.mean()
                features[f"rolling_std_{window}"] = recent.std()
                features[f"rolling_min_{window}"] = recent.min()
                features[f"rolling_max_{window}"] = recent.max()
            else:
                features[f"rolling_mean_{window}"] = np.nan
                features[f"rolling_std_{window}"] = np.nan
                features[f"rolling_min_{window}"] = np.nan
                features[f"rolling_max_{window}"] = np.nan

        # Trend feature
        features["trend"] = len(history)

        # Ensure column order matches training
        df = pd.DataFrame([features])
        df = df[self.feature_names_]

        return df

    def _calculate_confidence_intervals(
        self,
        forecast: pd.Series,
        horizon: int,
        alpha: float,
    ) -> pd.DataFrame:
        """Calculate prediction intervals using residual bootstrap."""
        # Estimate uncertainty from training residuals
        X_train, y_train = self._generate_training_features(self.training_series_)
        predictions = self.model_.predict(X_train)
        residuals = y_train.values - predictions
        residual_std = np.std(residuals)

        # Increase uncertainty with horizon
        z = stats.norm.ppf(1 - alpha / 2)
        h = np.arange(1, horizon + 1)
        uncertainty = residual_std * np.sqrt(1 + 0.1 * h)

        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_feature_importance(self, top_n: int = 20) -> pd.Series:
        """
        Get top feature importances.

        Parameters
        ----------
        top_n : int, default 20
            Number of top features to return.

        Returns
        -------
        importance : pd.Series
            Feature importances sorted descending.
        """
        self._check_is_fitted()
        return self.feature_importance_.head(top_n)

    def plot_feature_importance(self, top_n: int = 20):
        """
        Plot feature importances.

        Parameters
        ----------
        top_n : int, default 20
            Number of top features to plot.

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

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

        importance = self.get_feature_importance(top_n)

        fig, ax = plt.subplots(figsize=(10, 8))
        importance.plot(kind="barh", ax=ax)
        ax.set_xlabel("Importance")
        ax.set_ylabel("Feature")
        ax.set_title(f"Top {top_n} Feature Importances")
        ax.invert_yaxis()

        plt.tight_layout()
        return fig


def generate_temporal_features(dates: pd.DatetimeIndex) -> pd.DataFrame:
    """
    Generate temporal features from dates.

    Parameters
    ----------
    dates : pd.DatetimeIndex
        Dates to generate features for.

    Returns
    -------
    features : pd.DataFrame
        Temporal features.
    """
    features = pd.DataFrame(index=dates)

    features["day_of_week"] = dates.dayofweek
    features["day_of_month"] = dates.day
    features["month"] = dates.month
    features["week_of_year"] = dates.isocalendar().week.values
    features["is_weekend"] = (dates.dayofweek >= 5).astype(int)
    features["quarter"] = dates.quarter
    features["day_of_year"] = dates.dayofyear
    features["year"] = dates.year

    return features


def generate_lag_features(
    y: pd.Series,
    lags: List[int],
) -> pd.DataFrame:
    """
    Generate lag features from a time series.

    Parameters
    ----------
    y : pd.Series
        Time series.
    lags : List[int]
        Lag periods.

    Returns
    -------
    features : pd.DataFrame
        Lag features.
    """
    features = pd.DataFrame(index=y.index)

    for lag in lags:
        features[f"lag_{lag}"] = y.shift(lag)

    return features


def generate_rolling_features(
    y: pd.Series,
    windows: List[int],
    statistics: List[str] = None,
) -> pd.DataFrame:
    """
    Generate rolling window features.

    Parameters
    ----------
    y : pd.Series
        Time series.
    windows : List[int]
        Window sizes.
    statistics : List[str], optional
        Statistics to calculate. Default: ['mean', 'std'].

    Returns
    -------
    features : pd.DataFrame
        Rolling features.
    """
    if statistics is None:
        statistics = ["mean", "std"]

    features = pd.DataFrame(index=y.index)

    for window in windows:
        rolling = y.rolling(window=window)

        for stat in statistics:
            if stat == "mean":
                features[f"rolling_mean_{window}"] = rolling.mean()
            elif stat == "std":
                features[f"rolling_std_{window}"] = rolling.std()
            elif stat == "min":
                features[f"rolling_min_{window}"] = rolling.min()
            elif stat == "max":
                features[f"rolling_max_{window}"] = rolling.max()
            elif stat == "median":
                features[f"rolling_median_{window}"] = rolling.median()

    return features
