"""Seasonality extraction and decomposition methods."""

from dataclasses import dataclass
from typing import Literal, Optional, Union, List, Tuple

import numpy as np
import pandas as pd
from scipy import stats
from statsmodels.tsa.seasonal import STL, MSTL, seasonal_decompose


@dataclass
class DecompositionResult:
    """
    Result of time series decomposition.

    Attributes
    ----------
    observed : pd.Series
        Original observed values.
    trend : pd.Series
        Trend component.
    seasonal : pd.Series or pd.DataFrame
        Seasonal component(s). DataFrame if multiple seasonal periods.
    residual : pd.Series
        Residual (remainder) component.
    model : str
        Decomposition model ('additive' or 'multiplicative').
    seasonal_periods : List[int]
        Seasonal period(s) used in decomposition.
    """

    observed: pd.Series
    trend: pd.Series
    seasonal: Union[pd.Series, pd.DataFrame]
    residual: pd.Series
    model: str
    seasonal_periods: List[int]

    def get_seasonal_indices(self, period: Optional[int] = None) -> pd.Series:
        """
        Get seasonal indices for a specific period.

        Parameters
        ----------
        period : int, optional
            Seasonal period to get indices for. If None and only one
            seasonal component exists, returns that.

        Returns
        -------
        indices : pd.Series
            Seasonal indices indexed by position within the period.
        """
        if isinstance(self.seasonal, pd.DataFrame):
            if period is None:
                raise ValueError(
                    "Must specify period when multiple seasonal components exist"
                )
            col_name = f"seasonal_{period}"
            if col_name not in self.seasonal.columns:
                raise ValueError(f"No seasonal component for period {period}")
            seasonal_values = self.seasonal[col_name]
        else:
            seasonal_values = self.seasonal

        # Calculate average seasonal index for each position
        period = period or self.seasonal_periods[0]
        n = len(seasonal_values)

        indices = np.zeros(period)
        counts = np.zeros(period)

        for i, val in enumerate(seasonal_values.values):
            pos = i % period
            if not np.isnan(val):
                indices[pos] += val
                counts[pos] += 1

        # Average the indices
        with np.errstate(divide="ignore", invalid="ignore"):
            indices = np.where(counts > 0, indices / counts, 0)

        return pd.Series(indices, index=range(period), name="seasonal_index")


def extract_seasonality(
    y: pd.Series,
    period: int,
    model: Literal["additive", "multiplicative"] = "additive",
    extrapolate_trend: bool = True,
) -> DecompositionResult:
    """
    Extract seasonality using classical decomposition.

    Uses moving average for trend extraction and ratio/difference method
    for seasonal component.

    Parameters
    ----------
    y : pd.Series
        Time series to decompose. Should have DatetimeIndex.
    period : int
        Seasonal period (e.g., 7 for weekly seasonality in daily data).
    model : {'additive', 'multiplicative'}, default 'additive'
        Type of decomposition model.
    extrapolate_trend : bool, default True
        Whether to extrapolate trend at edges using linear regression.

    Returns
    -------
    result : DecompositionResult
        Decomposition result with trend, seasonal, and residual components.

    Examples
    --------
    >>> 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)
    >>> result = extract_seasonality(y, period=7, model='additive')
    >>> print(result.seasonal_periods)
    [7]
    """
    if not isinstance(y.index, pd.DatetimeIndex):
        raise ValueError("Series must have DatetimeIndex")

    if len(y) < 2 * period:
        raise ValueError(f"Need at least {2 * period} observations for period={period}")

    # Use statsmodels classical decomposition
    decomposition = seasonal_decompose(
        y,
        model=model,
        period=period,
        extrapolate_trend="freq" if extrapolate_trend else None,
    )

    return DecompositionResult(
        observed=y,
        trend=decomposition.trend,
        seasonal=decomposition.seasonal,
        residual=decomposition.resid,
        model=model,
        seasonal_periods=[period],
    )


def stl_decomposition(
    y: pd.Series,
    period: int,
    seasonal: int = 7,
    trend: Optional[int] = None,
    low_pass: Optional[int] = None,
    robust: bool = True,
) -> DecompositionResult:
    """
    Decompose time series using STL (Seasonal-Trend decomposition using LOESS).

    STL is more robust to outliers and handles missing values better than
    classical decomposition.

    Parameters
    ----------
    y : pd.Series
        Time series to decompose. Should have DatetimeIndex.
    period : int
        Seasonal period (e.g., 7 for weekly, 365 for annual).
    seasonal : int, default 7
        Seasonal smoothing parameter. Must be odd. Larger values give
        smoother seasonal component.
    trend : int, optional
        Trend smoothing parameter. Must be odd and >= 1.5 * period / (1 - 1.5/seasonal).
        If None, uses a good default based on period.
    low_pass : int, optional
        Low-pass filter length. Must be odd and >= period.
        If None, uses period + 1 (next odd).
    robust : bool, default True
        Whether to use robust fitting (iterative reweighting to reduce
        influence of outliers).

    Returns
    -------
    result : DecompositionResult
        Decomposition result with trend, seasonal, and residual components.

    Notes
    -----
    STL always produces an additive decomposition. For multiplicative data,
    apply log transform before STL and inverse transform after.

    References
    ----------
    Cleveland, R. B., Cleveland, W. S., McRae, J. E., & Terpenning, I. (1990).
    STL: A seasonal-trend decomposition procedure based on loess.
    Journal of Official Statistics, 6(1), 3-73.

    Examples
    --------
    >>> 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)
    >>> result = stl_decomposition(y, period=7, robust=True)
    """
    if not isinstance(y.index, pd.DatetimeIndex):
        raise ValueError("Series must have DatetimeIndex")

    if len(y) < 2 * period:
        raise ValueError(f"Need at least {2 * period} observations for period={period}")

    # Handle missing values by interpolation (STL requires complete data)
    y_filled = y.interpolate(method="linear").ffill().bfill()

    # Run STL decomposition
    stl = STL(
        y_filled,
        period=period,
        seasonal=seasonal,
        trend=trend,
        low_pass=low_pass,
        robust=robust,
    )
    decomposition = stl.fit()

    return DecompositionResult(
        observed=y,
        trend=pd.Series(decomposition.trend, index=y.index, name="trend"),
        seasonal=pd.Series(decomposition.seasonal, index=y.index, name="seasonal"),
        residual=pd.Series(decomposition.resid, index=y.index, name="residual"),
        model="additive",
        seasonal_periods=[period],
    )


def mstl_decomposition(
    y: pd.Series,
    periods: List[int],
    stl_kwargs: Optional[dict] = None,
) -> DecompositionResult:
    """
    Decompose time series with multiple seasonal periods using MSTL.

    MSTL (Multiple Seasonal-Trend decomposition using LOESS) handles
    multiple overlapping seasonal patterns, common in daily data with
    both weekly (7) and annual (365) seasonality.

    Parameters
    ----------
    y : pd.Series
        Time series to decompose. Should have DatetimeIndex.
    periods : List[int]
        List of seasonal periods in ascending order.
        Common examples:
        - Daily data: [7, 365] (weekly + annual)
        - Hourly data: [24, 168, 8766] (daily + weekly + annual)
    stl_kwargs : dict, optional
        Additional keyword arguments passed to the underlying STL.
        For example: {'robust': True, 'seasonal': 11}

    Returns
    -------
    result : DecompositionResult
        Decomposition result. The seasonal attribute is a DataFrame
        with columns 'seasonal_7', 'seasonal_365', etc.

    Notes
    -----
    MSTL decomposes in order of period length, extracting each seasonal
    component successively from the residuals of the previous extraction.

    References
    ----------
    Bandara, K., Hyndman, R. J., & Bergmeir, C. (2021).
    MSTL: A Seasonal-Trend Decomposition Algorithm for Time Series with
    Multiple Seasonal Patterns.

    Examples
    --------
    >>> dates = pd.date_range('2024-01-01', periods=730, freq='D')
    >>> weekly = 10 * np.sin(np.arange(730) * 2 * np.pi / 7)
    >>> annual = 20 * np.sin(np.arange(730) * 2 * np.pi / 365)
    >>> y = pd.Series(100 + weekly + annual, index=dates)
    >>> result = mstl_decomposition(y, periods=[7, 365])
    >>> print(result.seasonal.columns)
    Index(['seasonal_7', 'seasonal_365'], dtype='object')
    """
    if not isinstance(y.index, pd.DatetimeIndex):
        raise ValueError("Series must have DatetimeIndex")

    # Sort periods in ascending order
    periods = sorted(periods)

    min_required = 2 * max(periods)
    if len(y) < min_required:
        raise ValueError(
            f"Need at least {min_required} observations for periods={periods}"
        )

    # Handle missing values
    y_filled = y.interpolate(method="linear").ffill().bfill()

    # Default STL kwargs
    kwargs = {"robust": True}
    if stl_kwargs is not None:
        kwargs.update(stl_kwargs)

    # Run MSTL decomposition
    mstl = MSTL(y_filled, periods=periods, stl_kwargs=kwargs)
    decomposition = mstl.fit()

    # Extract seasonal components into DataFrame
    seasonal_df = pd.DataFrame(index=y.index)
    for i, period in enumerate(periods):
        col_name = f"seasonal_{period}"
        seasonal_df[col_name] = decomposition.seasonal.iloc[:, i].values

    return DecompositionResult(
        observed=y,
        trend=pd.Series(decomposition.trend, index=y.index, name="trend"),
        seasonal=seasonal_df,
        residual=pd.Series(decomposition.resid, index=y.index, name="residual"),
        model="additive",
        seasonal_periods=periods,
    )


def apply_seasonality(
    trend_forecast: pd.Series,
    seasonal_indices: Union[pd.Series, List[float]],
    period: int,
    model: Literal["additive", "multiplicative"] = "additive",
    start_position: int = 0,
) -> pd.Series:
    """
    Apply seasonal indices to a trend forecast.

    Parameters
    ----------
    trend_forecast : pd.Series
        Trend forecast values.
    seasonal_indices : pd.Series or list
        Seasonal indices for one complete period. Should have length
        equal to `period`.
    period : int
        Seasonal period.
    model : {'additive', 'multiplicative'}, default 'additive'
        How to apply seasonality:
        - 'additive': Y = T + S
        - 'multiplicative': Y = T * S
    start_position : int, default 0
        Position within the seasonal cycle for the first forecast.
        For example, if forecasting starts on Wednesday with weekly
        seasonality (Monday=0), use start_position=2.

    Returns
    -------
    forecast : pd.Series
        Seasonally adjusted forecast.

    Examples
    --------
    >>> trend = pd.Series([100, 101, 102, 103, 104, 105, 106])
    >>> indices = [0.9, 1.0, 1.1, 1.2, 1.1, 1.0, 0.7]  # Weekly pattern
    >>> forecast = apply_seasonality(trend, indices, period=7, model='multiplicative')
    """
    indices = np.array(seasonal_indices)

    if len(indices) != period:
        raise ValueError(f"seasonal_indices length ({len(indices)}) != period ({period})")

    # Create array of seasonal values aligned with forecast
    n = len(trend_forecast)
    seasonal_values = np.zeros(n)

    for i in range(n):
        pos = (start_position + i) % period
        seasonal_values[i] = indices[pos]

    # Apply seasonality
    if model == "additive":
        result = trend_forecast + seasonal_values
    elif model == "multiplicative":
        result = trend_forecast * seasonal_values
    else:
        raise ValueError(f"Unknown model: {model}")

    return pd.Series(result, index=trend_forecast.index, name="forecast")


def apply_multi_seasonality(
    trend_forecast: pd.Series,
    decomposition_result: DecompositionResult,
    forecast_dates: pd.DatetimeIndex,
    training_dates: pd.DatetimeIndex,
) -> pd.Series:
    """
    Apply multiple seasonal components to a trend forecast.

    Extrapolates seasonal patterns from training data to forecast dates.

    Parameters
    ----------
    trend_forecast : pd.Series
        Trend forecast values.
    decomposition_result : DecompositionResult
        Result from mstl_decomposition containing seasonal components.
    forecast_dates : pd.DatetimeIndex
        Dates for which to generate forecasts.
    training_dates : pd.DatetimeIndex
        Dates from the training data (used to determine seasonal phase).

    Returns
    -------
    forecast : pd.Series
        Seasonally adjusted forecast.
    """
    if decomposition_result.model != "additive":
        raise ValueError("Multi-seasonality only supports additive model")

    seasonal_total = np.zeros(len(forecast_dates))

    for period in decomposition_result.seasonal_periods:
        # Get seasonal indices for this period
        indices = decomposition_result.get_seasonal_indices(period)

        # Determine starting position based on last training date
        last_train_date = training_dates[-1]
        days_since_reference = (forecast_dates[0] - last_train_date).days

        # For daily data, use day-based positioning
        # The last training observation corresponds to position (len(training_dates)-1) % period
        last_position = (len(training_dates) - 1) % period
        start_position = (last_position + days_since_reference) % period

        # Apply this seasonal component
        for i in range(len(forecast_dates)):
            pos = (start_position + i) % period
            seasonal_total[i] += indices.iloc[pos]

    # Additive: Y = T + S1 + S2 + ...
    forecast = trend_forecast.values + seasonal_total

    return pd.Series(forecast, index=forecast_dates, name="forecast")


def estimate_seasonal_strength(
    decomposition: DecompositionResult,
    period: Optional[int] = None,
) -> float:
    """
    Estimate the strength of seasonality in a decomposition.

    Strength ranges from 0 (no seasonality) to 1 (purely seasonal).

    Parameters
    ----------
    decomposition : DecompositionResult
        Result from decomposition.
    period : int, optional
        Which seasonal period to evaluate (for MSTL results).

    Returns
    -------
    strength : float
        Seasonal strength between 0 and 1.

    Notes
    -----
    Strength is calculated as: 1 - Var(R) / Var(S + R)

    where S is seasonal component and R is residual.

    References
    ----------
    Wang, X., Smith, K. A., & Hyndman, R. J. (2006).
    Characteristic-based clustering for time series data.
    """
    if isinstance(decomposition.seasonal, pd.DataFrame):
        if period is None:
            # Sum all seasonal components
            seasonal = decomposition.seasonal.sum(axis=1)
        else:
            col = f"seasonal_{period}"
            if col not in decomposition.seasonal.columns:
                raise ValueError(f"No seasonal component for period {period}")
            seasonal = decomposition.seasonal[col]
    else:
        seasonal = decomposition.seasonal

    residual = decomposition.residual

    # Remove NaN values
    mask = ~(seasonal.isna() | residual.isna())
    seasonal = seasonal[mask]
    residual = residual[mask]

    var_residual = residual.var()
    var_seasonal_plus_residual = (seasonal + residual).var()

    if var_seasonal_plus_residual == 0:
        return 0.0

    strength = max(0, 1 - var_residual / var_seasonal_plus_residual)
    return strength


def estimate_trend_strength(decomposition: DecompositionResult) -> float:
    """
    Estimate the strength of trend in a decomposition.

    Strength ranges from 0 (no trend) to 1 (purely trend).

    Parameters
    ----------
    decomposition : DecompositionResult
        Result from decomposition.

    Returns
    -------
    strength : float
        Trend strength between 0 and 1.

    Notes
    -----
    Strength is calculated as: 1 - Var(R) / Var(T + R)

    where T is trend component and R is residual.
    """
    trend = decomposition.trend
    residual = decomposition.residual

    # Remove NaN values
    mask = ~(trend.isna() | residual.isna())
    trend = trend[mask]
    residual = residual[mask]

    var_residual = residual.var()
    var_trend_plus_residual = (trend + residual).var()

    if var_trend_plus_residual == 0:
        return 0.0

    strength = max(0, 1 - var_residual / var_trend_plus_residual)
    return strength


def detect_seasonality(
    y: pd.Series,
    max_period: int = 365,
    significance: float = 0.05,
) -> List[int]:
    """
    Automatically detect seasonal periods in a time series.

    Uses autocorrelation analysis to identify significant seasonal periods.

    Parameters
    ----------
    y : pd.Series
        Time series to analyze.
    max_period : int, default 365
        Maximum period to check.
    significance : float, default 0.05
        Significance level for detecting peaks.

    Returns
    -------
    periods : List[int]
        Detected seasonal periods in ascending order.

    Examples
    --------
    >>> dates = pd.date_range('2024-01-01', periods=365, freq='D')
    >>> weekly = 10 * np.sin(np.arange(365) * 2 * np.pi / 7)
    >>> y = pd.Series(100 + weekly, index=dates)
    >>> periods = detect_seasonality(y)
    >>> print(7 in periods)
    True
    """
    n = len(y)
    max_period = min(max_period, n // 2)

    # Calculate autocorrelation
    y_centered = y - y.mean()
    acf = np.correlate(y_centered, y_centered, mode="full")
    acf = acf[n - 1 :]  # Take positive lags only
    acf = acf / acf[0]  # Normalize

    # Calculate significance threshold
    # Under null hypothesis (white noise), ACF ~ N(0, 1/n)
    threshold = stats.norm.ppf(1 - significance / 2) / np.sqrt(n)

    # Find peaks in ACF
    periods = []

    for lag in range(2, max_period + 1):
        if lag >= len(acf):
            break

        # Check if this lag is a local maximum and significant
        is_peak = (
            acf[lag] > acf[lag - 1]
            and (lag + 1 >= len(acf) or acf[lag] > acf[lag + 1])
        )
        is_significant = acf[lag] > threshold

        if is_peak and is_significant:
            # Additional check: is this a primary period or harmonic?
            is_harmonic = False
            for existing_period in periods:
                if lag % existing_period == 0 or existing_period % lag == 0:
                    # Could be a harmonic
                    if acf[lag] < acf[existing_period % len(acf)] * 0.8:
                        is_harmonic = True
                        break

            if not is_harmonic:
                periods.append(lag)

    return sorted(periods)


def select_decomposition_model(
    y: pd.Series,
    period: int,
) -> str:
    """
    Select between additive and multiplicative decomposition.

    Multiplicative is preferred when seasonal variation increases
    proportionally with the level of the series.

    Parameters
    ----------
    y : pd.Series
        Time series to analyze.
    period : int
        Seasonal period.

    Returns
    -------
    model : str
        'additive' or 'multiplicative'.

    Notes
    -----
    Uses coefficient of variation analysis: if CV increases with level,
    multiplicative is preferred.
    """
    # Check for zeros or negative values (can't use multiplicative)
    if (y <= 0).any():
        return "additive"

    # Split series into seasonal blocks
    n = len(y)
    n_blocks = n // period

    if n_blocks < 3:
        return "additive"

    # Calculate mean and std for each block
    means = []
    stds = []

    for i in range(n_blocks):
        block = y.iloc[i * period : (i + 1) * period]
        means.append(block.mean())
        stds.append(block.std())

    means = np.array(means)
    stds = np.array(stds)

    # Check correlation between mean and std
    if len(means) > 2 and np.std(means) > 0 and np.std(stds) > 0:
        correlation = np.corrcoef(means, stds)[0, 1]

        # If positive correlation, seasonal variation increases with level
        if correlation > 0.5:
            return "multiplicative"

    return "additive"
