"""Analog (Copy) model for event-driven forecasting."""

from typing import Optional, Union, Tuple, List, Literal
from datetime import datetime

import numpy as np
import pandas as pd

from ad_inventory_forecast.core.base import BaseForecaster
from ad_inventory_forecast.models.trend import (
    BaseTrend,
    LinearTrend,
    select_best_trend,
)
from ad_inventory_forecast.utils.dtw import (
    align_sequences,
    scale_and_warp,
    EventPatternLibrary,
)


class AnalogEventForecaster(BaseForecaster):
    """
    Analog (Copy) model for forecasting event-driven traffic spikes.

    The Copy model forecasts future event traffic by scaling historical
    event patterns based on trend growth. It's ideal for predictable
    recurring events like Super Bowl, Black Friday, or seasonal campaigns.

    The core formula is:
        Y_forecast = Y_source × G

    where G = T_target / T_source is the Growth Factor derived from
    trend analysis.

    Parameters
    ----------
    source_window : Tuple[str, str]
        Date range (start, end) of the historical analog event.
        Example: ('2023-02-12', '2023-02-13') for Super Bowl 2023.
    target_window : Tuple[str, str]
        Date range (start, end) of the future event to forecast.
        Example: ('2024-02-11', '2024-02-12') for Super Bowl 2024.
    use_dtw : bool, default True
        Whether to use Dynamic Time Warping for alignment when source
        and target events have different durations.
    alignment_method : {'dtw', 'linear', 'nearest'}, default 'dtw'
        Method for aligning sequences of different lengths.
    trend_model : str or BaseTrend, default 'auto'
        Trend model to use for growth factor calculation.
        Options: 'linear', 'exponential', 'auto', or a BaseTrend instance.
    scale_by_trend : bool, default True
        Whether to apply trend-based growth factor.
    confidence_level : float, default 0.95
        Confidence level for prediction intervals.

    Attributes
    ----------
    is_fitted_ : bool
        Whether the model has been fitted.
    source_values_ : np.ndarray
        Values from the source (historical) event.
    trend_model_ : BaseTrend
        Fitted trend model.
    growth_factor_ : float
        Calculated growth factor G = T_target / T_source.
    trend_at_source_ : float
        Trend value at source event.
    trend_at_target_ : float
        Trend value at target event.

    Examples
    --------
    >>> import pandas as pd
    >>> import numpy as np
    >>> # Create sample data with yearly trend and event spike
    >>> dates = pd.date_range('2022-01-01', periods=730, freq='D')
    >>> trend = np.linspace(100, 120, 730)
    >>> y = pd.Series(trend + np.random.randn(730) * 5, index=dates)
    >>> # Add event spike (e.g., Super Bowl)
    >>> y.loc['2023-02-12':'2023-02-13'] *= 3
    >>>
    >>> # Forecast next year's event
    >>> forecaster = AnalogEventForecaster(
    ...     source_window=('2023-02-12', '2023-02-13'),
    ...     target_window=('2024-02-11', '2024-02-12'),
    ... )
    >>> forecaster.fit(y)
    >>> forecast = forecaster.predict()
    >>> print(f"Growth factor: {forecaster.growth_factor_:.2f}")

    Notes
    -----
    Best practices for the Copy model:
    1. Use events with similar characteristics (same sport, same type)
    2. Ensure sufficient pre-event history for trend estimation
    3. Consider using multiple analog events and averaging
    4. Review the growth factor for reasonableness
    """

    def __init__(
        self,
        source_window: Tuple[str, str],
        target_window: Tuple[str, str],
        use_dtw: bool = True,
        alignment_method: Literal["dtw", "linear", "nearest"] = "dtw",
        trend_model: Union[str, BaseTrend] = "auto",
        scale_by_trend: bool = True,
        confidence_level: float = 0.95,
    ):
        """Initialize the AnalogEventForecaster."""
        super().__init__()

        self.source_window = source_window
        self.target_window = target_window
        self.use_dtw = use_dtw
        self.alignment_method = alignment_method
        self.trend_model = trend_model
        self.scale_by_trend = scale_by_trend
        self.confidence_level = confidence_level

        # Fitted attributes
        self.source_values_ = None
        self.trend_model_ = None
        self.growth_factor_ = None
        self.trend_at_source_ = None
        self.trend_at_target_ = None
        self._source_start_idx = None
        self._source_end_idx = None
        self._forecast_values = None

    def fit(
        self,
        y: Union[pd.Series, np.ndarray],
        X: Optional[Union[pd.DataFrame, np.ndarray]] = None,
    ) -> "AnalogEventForecaster":
        """
        Fit the analog forecaster to historical data.

        Parameters
        ----------
        y : pd.Series or np.ndarray
            Historical time series. Must include the source event window.
        X : pd.DataFrame or np.ndarray, optional
            Exogenous features (not used).

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

        # Parse window dates
        source_start = pd.Timestamp(self.source_window[0])
        source_end = pd.Timestamp(self.source_window[1])
        target_start = pd.Timestamp(self.target_window[0])
        target_end = pd.Timestamp(self.target_window[1])

        # Validate source window is in data
        if source_start < y.index.min() or source_end > y.index.max():
            data_start = y.index.min().strftime('%Y-%m-%d')
            data_end = y.index.max().strftime('%Y-%m-%d')
            source_start_str = source_start.strftime('%Y-%m-%d')
            source_end_str = source_end.strftime('%Y-%m-%d')

            # Calculate how much more history is needed
            if source_start < y.index.min():
                days_needed = (y.index.min() - source_start).days
                suggestion = (
                    f"You need at least {days_needed} more days of historical data "
                    f"before {data_start} to include the analog event period."
                )
            else:
                days_needed = (source_end - y.index.max()).days
                suggestion = (
                    f"The source window extends {days_needed} days beyond your data. "
                    f"Ensure your data includes the complete analog event period."
                )

            raise ValueError(
                f"Analog source window not available in training data.\n"
                f"\n"
                f"  Source window requested: {source_start_str} to {source_end_str}\n"
                f"  Training data available: {data_start} to {data_end}\n"
                f"\n"
                f"  {suggestion}\n"
                f"\n"
                f"  Alternatives:\n"
                f"  - Use a different analog event that falls within your data range\n"
                f"  - Extend your historical data to include the analog period\n"
                f"  - Consider using InventoryForecaster for baseline forecasting"
            )

        # Extract source event values
        source_mask = (y.index >= source_start) & (y.index <= source_end)
        self.source_values_ = y[source_mask].values

        if len(self.source_values_) == 0:
            raise ValueError(f"No data found in source window {self.source_window}")

        # Store source indices
        self._source_start_idx = y.index.get_indexer([source_start], method="nearest")[0]
        self._source_end_idx = y.index.get_indexer([source_end], method="nearest")[0]

        # Fit trend model on non-event data
        self.trend_model_ = self._fit_trend_model(y, source_mask)

        # Calculate growth factor
        if self.scale_by_trend:
            self.growth_factor_ = self._calculate_growth_factor(
                y, source_start, source_end, target_start, target_end
            )
        else:
            self.growth_factor_ = 1.0

        self.is_fitted_ = True
        return self

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

        Parameters
        ----------
        horizon : int, optional
            Not used (horizon is determined by target_window).
            Kept for API compatibility.
        return_conf_int : bool, default False
            Whether to return confidence intervals (legacy parameter).
        alpha : float, default 0.05
            Significance level for confidence intervals.
        return_df : bool, default True
            If True (default), returns DataFrame with 'predicted_impressions'
            and 'event_label' columns (SDK-compliant format).
        event_label : str, optional
            Label for the event. If None, generates from target window.

        Returns
        -------
        forecast_df : pd.DataFrame
            If return_df=True (default): DataFrame with columns:
            - 'predicted_impressions': Point forecasts
            - 'event_label': Label for each forecast day
        forecasts : pd.Series
            If return_df=False: Point forecasts for the target event window.
        conf_int : pd.DataFrame, optional
            Confidence intervals. Only returned if return_conf_int=True.
        """
        self._check_is_fitted()

        # Parse target window
        target_start = pd.Timestamp(self.target_window[0])
        target_end = pd.Timestamp(self.target_window[1])

        # Generate forecast dates
        target_dates = pd.date_range(
            start=target_start,
            end=target_end,
            freq=self._get_freq_code(),
        )
        target_length = len(target_dates)

        # Align source to target duration
        if len(self.source_values_) != target_length:
            if self.use_dtw:
                aligned_values = align_sequences(
                    self.source_values_,
                    target_length,
                    method=self.alignment_method,
                )
            else:
                # Simple linear interpolation
                aligned_values = align_sequences(
                    self.source_values_,
                    target_length,
                    method="linear",
                )
        else:
            aligned_values = self.source_values_.copy()

        # Apply growth factor
        forecast_values = aligned_values * self.growth_factor_

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

        self._forecast_values = forecast_values

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

        # Return SDK-compliant DataFrame by default
        if return_df and not return_conf_int:
            # Generate event labels
            if event_label is None:
                base_label = f"Event_{target_start.strftime('%Y%m%d')}"
            else:
                base_label = event_label

            event_labels = [f"{base_label}_Day{i+1}" for i in range(target_length)]

            return pd.DataFrame({
                "predicted_impressions": forecast_values,
                "event_label": event_labels,
            }, index=target_dates)

        if return_conf_int:
            conf_int = self._calculate_confidence_intervals(
                forecast, alpha
            )
            return forecast, conf_int

        return forecast

    def _fit_trend_model(
        self,
        y: pd.Series,
        event_mask: np.ndarray,
    ) -> BaseTrend:
        """Fit trend model excluding event period."""
        # Exclude event from trend fitting
        y_non_event = y[~event_mask]
        t_non_event = np.where(~event_mask)[0]

        if len(y_non_event) < 30:
            # Not enough data, use all data
            t = np.arange(len(y))
            values = y.values
        else:
            t = t_non_event
            values = y_non_event.values

        if isinstance(self.trend_model, BaseTrend):
            trend = self.trend_model
            trend.fit(values, t)
        elif self.trend_model == "auto":
            trend = select_best_trend(values, t)
        elif self.trend_model == "linear":
            trend = LinearTrend()
            trend.fit(values, t)
        else:
            from ad_inventory_forecast.models.trend import TREND_MODELS
            if self.trend_model in TREND_MODELS:
                trend = TREND_MODELS[self.trend_model]()
                trend.fit(values, t)
            else:
                raise ValueError(f"Unknown trend model: {self.trend_model}")

        return trend

    def _calculate_growth_factor(
        self,
        y: pd.Series,
        source_start: pd.Timestamp,
        source_end: pd.Timestamp,
        target_start: pd.Timestamp,
        target_end: pd.Timestamp,
    ) -> float:
        """Calculate growth factor G = T_target / T_source."""
        # Get midpoints
        source_mid_idx = (self._source_start_idx + self._source_end_idx) // 2

        # Calculate time difference
        days_diff = (target_start - source_start).days
        freq_days = {"D": 1, "W": 7, "M": 30}.get(self.freq_, 1)
        periods_diff = days_diff / freq_days

        target_mid_idx = source_mid_idx + int(periods_diff)

        # Get trend values
        self.trend_at_source_ = self.trend_model_.predict(np.array([source_mid_idx]))[0]
        self.trend_at_target_ = self.trend_model_.predict(np.array([target_mid_idx]))[0]

        # Avoid division by zero
        if self.trend_at_source_ <= 0:
            return 1.0

        growth_factor = self.trend_at_target_ / self.trend_at_source_

        # Cap extreme growth factors
        growth_factor = np.clip(growth_factor, 0.5, 2.0)

        return growth_factor

    def _calculate_confidence_intervals(
        self,
        forecast: pd.Series,
        alpha: float,
    ) -> pd.DataFrame:
        """Calculate prediction intervals based on historical variance."""
        from scipy import stats

        # Estimate variance from source event residuals
        # (difference between actual and scaled trend)
        source_trend = self.trend_model_.predict(
            np.arange(self._source_start_idx, self._source_end_idx + 1)
        )

        if len(source_trend) == len(self.source_values_):
            residuals = self.source_values_ - source_trend
            residual_std = np.std(residuals)
        else:
            # Use coefficient of variation from source values
            cv = np.std(self.source_values_) / np.mean(self.source_values_)
            residual_std = forecast.values * cv

        if np.isscalar(residual_std):
            residual_std = np.full(len(forecast), residual_std)

        # Scale residual std by growth factor
        residual_std = residual_std * self.growth_factor_

        # Calculate intervals
        z = stats.norm.ppf(1 - alpha / 2)
        lower = forecast.values - z * residual_std
        upper = forecast.values + z * residual_std

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

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

    def _get_freq_code(self) -> str:
        """Get pandas frequency code."""
        return {"D": "D", "W": "W", "M": "MS"}.get(self.freq_, "D")

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

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

        return {
            "source_window": self.source_window,
            "target_window": self.target_window,
            "source_duration": len(self.source_values_),
            "source_mean": np.mean(self.source_values_),
            "source_max": np.max(self.source_values_),
            "growth_factor": self.growth_factor_,
            "trend_at_source": self.trend_at_source_,
            "trend_at_target": self.trend_at_target_,
            "trend_model": type(self.trend_model_).__name__,
            "use_dtw": self.use_dtw,
        }

    def plot_comparison(self, forecast: Optional[pd.Series] = None):
        """
        Plot source event vs forecasted event.

        Parameters
        ----------
        forecast : pd.Series, optional
            Forecast to plot. If None, generates new forecast.

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

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

        if forecast is None:
            forecast = self.predict()

        fig, axes = plt.subplots(1, 2, figsize=(14, 5))

        # Plot source event
        source_dates = self.training_series_.index[
            self._source_start_idx : self._source_end_idx + 1
        ]
        axes[0].plot(
            range(len(self.source_values_)),
            self.source_values_,
            "b-o",
            label="Source Event",
            markersize=4,
        )
        axes[0].set_title(f"Source Event\n{self.source_window[0]} to {self.source_window[1]}")
        axes[0].set_xlabel("Day")
        axes[0].set_ylabel("Impressions")
        axes[0].legend()

        # Plot forecast
        axes[1].plot(
            range(len(forecast)),
            forecast.values,
            "r-o",
            label="Forecast",
            markersize=4,
        )
        axes[1].set_title(
            f"Target Event Forecast\n{self.target_window[0]} to {self.target_window[1]}\n"
            f"(Growth Factor: {self.growth_factor_:.2f})"
        )
        axes[1].set_xlabel("Day")
        axes[1].set_ylabel("Impressions")
        axes[1].legend()

        plt.tight_layout()
        return fig


class MultiAnalogForecaster(BaseForecaster):
    """
    Forecaster that combines multiple analog events.

    Uses multiple historical events as analogs and combines their
    forecasts (weighted average or ensemble).

    Parameters
    ----------
    source_windows : List[Tuple[str, str]]
        List of date ranges for historical analog events.
    target_window : Tuple[str, str]
        Date range for the future event to forecast.
    weights : List[float], optional
        Weights for each analog event. If None, uses equal weights.
        More recent events can be given higher weights.
    combination_method : {'weighted_mean', 'median', 'weighted_median'}, default 'weighted_mean'
        Method for combining forecasts from multiple analogs.

    Examples
    --------
    >>> forecaster = MultiAnalogForecaster(
    ...     source_windows=[
    ...         ('2022-02-13', '2022-02-14'),  # Super Bowl 2022
    ...         ('2023-02-12', '2023-02-13'),  # Super Bowl 2023
    ...     ],
    ...     target_window=('2024-02-11', '2024-02-12'),
    ...     weights=[0.3, 0.7],  # More weight on recent event
    ... )
    >>> forecaster.fit(y)
    >>> forecast = forecaster.predict()
    """

    def __init__(
        self,
        source_windows: List[Tuple[str, str]],
        target_window: Tuple[str, str],
        weights: Optional[List[float]] = None,
        combination_method: Literal[
            "weighted_mean", "median", "weighted_median"
        ] = "weighted_mean",
        **kwargs,
    ):
        """Initialize the MultiAnalogForecaster."""
        super().__init__()

        self.source_windows = source_windows
        self.target_window = target_window
        self.weights = weights
        self.combination_method = combination_method
        self.kwargs = kwargs

        # Fitted attributes
        self.forecasters_ = []
        self.individual_forecasts_ = []

    def fit(
        self,
        y: Union[pd.Series, np.ndarray],
        X: Optional[Union[pd.DataFrame, np.ndarray]] = None,
    ) -> "MultiAnalogForecaster":
        """Fit individual analog forecasters for each source event."""
        y = self._validate_y(y)
        self.training_series_ = y.copy()
        self.n_obs_ = len(y)
        self.freq_ = self._infer_frequency(y)

        self.forecasters_ = []

        for source_window in self.source_windows:
            forecaster = AnalogEventForecaster(
                source_window=source_window,
                target_window=self.target_window,
                **self.kwargs,
            )
            try:
                forecaster.fit(y)
                self.forecasters_.append(forecaster)
            except ValueError as e:
                # Skip invalid source windows
                print(f"Warning: Skipping source {source_window}: {e}")

        if not self.forecasters_:
            raise ValueError("No valid source windows found")

        self.is_fitted_ = True
        return self

    def predict(
        self,
        horizon: Optional[int] = None,
        return_conf_int: bool = False,
        alpha: float = 0.05,
        return_df: bool = True,
        event_label: Optional[str] = None,
    ) -> Union[pd.Series, pd.DataFrame, Tuple[pd.Series, pd.DataFrame]]:
        """Generate combined forecast from multiple analogs.

        Parameters
        ----------
        horizon : int, optional
            Not used (kept for API compatibility).
        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'
            and 'event_label' columns (SDK-compliant format).
        event_label : str, optional
            Label for the event.

        Returns
        -------
        forecast_df : pd.DataFrame
            If return_df=True (default): DataFrame with prediction columns.
        forecasts : pd.Series
            If return_df=False: Point forecasts.
        conf_int : pd.DataFrame, optional
            Confidence intervals. Only returned if return_conf_int=True.
        """
        self._check_is_fitted()

        # Get individual forecasts (use return_df=False to get Series)
        forecasts = []
        forecast_index = None
        for forecaster in self.forecasters_:
            forecast_df = forecaster.predict(return_df=True)
            forecasts.append(forecast_df["predicted_impressions"].values)
            if forecast_index is None:
                forecast_index = forecast_df.index

        self.individual_forecasts_ = forecasts

        # Determine weights
        if self.weights is not None:
            weights = self.weights[: len(forecasts)]
        else:
            weights = [1.0] * len(forecasts)

        # Normalize weights
        weights = np.array(weights)
        weights = weights / weights.sum()

        # Combine forecasts
        forecasts_array = np.array(forecasts)

        if self.combination_method == "weighted_mean":
            combined = np.average(forecasts_array, axis=0, weights=weights)
        elif self.combination_method == "median":
            combined = np.median(forecasts_array, axis=0)
        elif self.combination_method == "weighted_median":
            # Approximate weighted median using repeated values
            combined = np.zeros(forecasts_array.shape[1])
            for i in range(forecasts_array.shape[1]):
                values = forecasts_array[:, i]
                # Create weighted samples
                n_samples = 1000
                samples = np.repeat(
                    values,
                    (weights * n_samples).astype(int),
                )
                combined[i] = np.median(samples)
        else:
            raise ValueError(f"Unknown combination method: {self.combination_method}")

        # Return SDK-compliant DataFrame by default
        if return_df and not return_conf_int:
            target_start = pd.Timestamp(self.target_window[0])
            if event_label is None:
                base_label = f"MultiEvent_{target_start.strftime('%Y%m%d')}"
            else:
                base_label = event_label

            event_labels = [f"{base_label}_Day{i+1}" for i in range(len(combined))]

            return pd.DataFrame({
                "predicted_impressions": combined,
                "event_label": event_labels,
            }, index=forecast_index)

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

        if return_conf_int:
            # Use spread of individual forecasts for uncertainty
            lower = np.min(forecasts_array, axis=0)
            upper = np.max(forecasts_array, axis=0)

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

        return forecast
