"""Fourier-based seasonality modeling."""

from typing import Optional, Union, List

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


class FourierSeasonality:
    """
    Fourier series seasonality model.

    Models seasonal patterns as a sum of sine and cosine waves:
    s(t) = Σ[aₙcos(2πnt/P) + bₙsin(2πnt/P)]

    This approach is flexible and can capture complex seasonal patterns
    with few parameters.

    Parameters
    ----------
    period : int
        Seasonal period (e.g., 7 for weekly, 365 for annual).
    n_terms : int, default 3
        Number of Fourier terms (harmonics) to use. Higher values
        capture more complex patterns but risk overfitting.
        Recommended: 3-5 for weekly, 5-10 for annual.

    Attributes
    ----------
    coefficients_ : np.ndarray
        Fitted Fourier coefficients [a1, b1, a2, b2, ..., an, bn].
    is_fitted_ : bool
        Whether the model has been fitted.

    Examples
    --------
    >>> dates = pd.date_range('2024-01-01', periods=365, freq='D')
    >>> seasonal = 10 * np.sin(np.arange(365) * 2 * np.pi / 7)
    >>> y = pd.Series(seasonal, index=dates)
    >>> fourier = FourierSeasonality(period=7, n_terms=3)
    >>> fourier.fit(y)
    >>> forecast_indices = fourier.predict(np.arange(30))
    """

    def __init__(self, period: int, n_terms: int = 3):
        """Initialize the Fourier seasonality model."""
        if period < 2:
            raise ValueError("Period must be >= 2")
        if n_terms < 1:
            raise ValueError("n_terms must be >= 1")
        if n_terms > period // 2:
            raise ValueError(f"n_terms ({n_terms}) cannot exceed period/2 ({period // 2})")

        self.period = period
        self.n_terms = n_terms
        self.coefficients_ = None
        self.is_fitted_ = False

    def _fourier_features(self, t: np.ndarray) -> np.ndarray:
        """
        Generate Fourier features for given time indices.

        Parameters
        ----------
        t : np.ndarray
            Time indices (can be float for fractional periods).

        Returns
        -------
        features : np.ndarray
            Feature matrix of shape (len(t), 2 * n_terms).
            Columns are [cos(2π*1*t/P), sin(2π*1*t/P), cos(2π*2*t/P), ...]
        """
        features = np.zeros((len(t), 2 * self.n_terms))

        for k in range(1, self.n_terms + 1):
            angle = 2 * np.pi * k * t / self.period
            features[:, 2 * (k - 1)] = np.cos(angle)
            features[:, 2 * (k - 1) + 1] = np.sin(angle)

        return features

    def fit(
        self,
        y: Union[pd.Series, np.ndarray],
        t: Optional[np.ndarray] = None,
    ) -> "FourierSeasonality":
        """
        Fit Fourier seasonality model to data.

        Parameters
        ----------
        y : pd.Series or np.ndarray
            Seasonal component or detrended series to fit.
        t : np.ndarray, optional
            Time indices corresponding to y. If None, uses 0, 1, 2, ...

        Returns
        -------
        self : FourierSeasonality
            Fitted model.
        """
        y = np.asarray(y)

        if t is None:
            t = np.arange(len(y))
        else:
            t = np.asarray(t)

        if len(t) != len(y):
            raise ValueError("t and y must have the same length")

        # Remove NaN values for fitting
        mask = ~np.isnan(y)
        y_valid = y[mask]
        t_valid = t[mask]

        if len(y_valid) < 2 * self.n_terms:
            raise ValueError(
                f"Need at least {2 * self.n_terms} non-NaN observations "
                f"for {self.n_terms} Fourier terms"
            )

        # Generate Fourier features
        X = self._fourier_features(t_valid)

        # Fit using least squares
        # Solve X @ coefficients = y
        self.coefficients_, _, _, _ = np.linalg.lstsq(X, y_valid, rcond=None)

        self.is_fitted_ = True
        return self

    def predict(self, t: Union[np.ndarray, int, List[int]]) -> np.ndarray:
        """
        Predict seasonal values for given time indices.

        Parameters
        ----------
        t : np.ndarray, int, or list
            Time indices to predict for.

        Returns
        -------
        seasonal : np.ndarray
            Predicted seasonal values.
        """
        if not self.is_fitted_:
            raise ValueError("Model must be fitted before predict()")

        t = np.atleast_1d(t).astype(float)
        X = self._fourier_features(t)

        return X @ self.coefficients_

    def get_seasonal_indices(self) -> pd.Series:
        """
        Get seasonal indices for one complete period.

        Returns
        -------
        indices : pd.Series
            Seasonal values for positions 0 to period-1.
        """
        if not self.is_fitted_:
            raise ValueError("Model must be fitted first")

        t = np.arange(self.period)
        values = self.predict(t)

        return pd.Series(values, index=t, name="seasonal_index")

    def get_amplitude(self, harmonic: int = 1) -> float:
        """
        Get amplitude of a specific harmonic.

        Parameters
        ----------
        harmonic : int, default 1
            Which harmonic (1 = fundamental frequency).

        Returns
        -------
        amplitude : float
            Amplitude of the harmonic.
        """
        if not self.is_fitted_:
            raise ValueError("Model must be fitted first")

        if harmonic < 1 or harmonic > self.n_terms:
            raise ValueError(f"harmonic must be between 1 and {self.n_terms}")

        a = self.coefficients_[2 * (harmonic - 1)]
        b = self.coefficients_[2 * (harmonic - 1) + 1]

        return np.sqrt(a**2 + b**2)

    def get_phase(self, harmonic: int = 1) -> float:
        """
        Get phase of a specific harmonic in radians.

        Parameters
        ----------
        harmonic : int, default 1
            Which harmonic (1 = fundamental frequency).

        Returns
        -------
        phase : float
            Phase in radians.
        """
        if not self.is_fitted_:
            raise ValueError("Model must be fitted first")

        if harmonic < 1 or harmonic > self.n_terms:
            raise ValueError(f"harmonic must be between 1 and {self.n_terms}")

        a = self.coefficients_[2 * (harmonic - 1)]
        b = self.coefficients_[2 * (harmonic - 1) + 1]

        return np.arctan2(b, a)

    def get_explained_variance(self, y: np.ndarray, t: Optional[np.ndarray] = None) -> float:
        """
        Calculate proportion of variance explained by the Fourier model.

        Parameters
        ----------
        y : np.ndarray
            Original values.
        t : np.ndarray, optional
            Time indices. If None, uses 0, 1, 2, ...

        Returns
        -------
        r_squared : float
            R-squared value (0 to 1).
        """
        if not self.is_fitted_:
            raise ValueError("Model must be fitted first")

        y = np.asarray(y)

        if t is None:
            t = np.arange(len(y))

        predicted = self.predict(t)

        # Remove NaN
        mask = ~np.isnan(y)
        y_valid = y[mask]
        pred_valid = predicted[mask]

        ss_res = np.sum((y_valid - pred_valid) ** 2)
        ss_tot = np.sum((y_valid - np.mean(y_valid)) ** 2)

        if ss_tot == 0:
            return 1.0 if ss_res == 0 else 0.0

        return 1 - ss_res / ss_tot


def generate_fourier_features(
    dates: pd.DatetimeIndex,
    period: int,
    n_terms: int = 3,
    reference_date: Optional[pd.Timestamp] = None,
) -> pd.DataFrame:
    """
    Generate Fourier feature columns for a DatetimeIndex.

    Useful for including seasonal features in regression models.

    Parameters
    ----------
    dates : pd.DatetimeIndex
        Dates to generate features for.
    period : int
        Seasonal period in days (e.g., 7 for weekly, 365 for annual).
    n_terms : int, default 3
        Number of Fourier terms.
    reference_date : pd.Timestamp, optional
        Reference date for calculating time index. If None, uses the
        minimum date in the series.

    Returns
    -------
    features : pd.DataFrame
        DataFrame with Fourier feature columns.
        Columns named: cos_1, sin_1, cos_2, sin_2, ...

    Examples
    --------
    >>> dates = pd.date_range('2024-01-01', periods=30, freq='D')
    >>> features = generate_fourier_features(dates, period=7, n_terms=3)
    >>> print(features.columns.tolist())
    ['cos_1', 'sin_1', 'cos_2', 'sin_2', 'cos_3', 'sin_3']
    """
    if reference_date is None:
        reference_date = dates.min()

    # Calculate time index (days since reference)
    t = (dates - reference_date).days.values.astype(float)

    features = pd.DataFrame(index=dates)

    for k in range(1, n_terms + 1):
        angle = 2 * np.pi * k * t / period
        features[f"cos_{k}"] = np.cos(angle)
        features[f"sin_{k}"] = np.sin(angle)

    return features


def select_optimal_n_terms(
    y: Union[pd.Series, np.ndarray],
    period: int,
    max_terms: Optional[int] = None,
    criterion: str = "aic",
) -> int:
    """
    Select optimal number of Fourier terms using information criterion.

    Parameters
    ----------
    y : pd.Series or np.ndarray
        Time series or seasonal component.
    period : int
        Seasonal period.
    max_terms : int, optional
        Maximum number of terms to consider. Default is period // 2.
    criterion : str, default 'aic'
        Information criterion: 'aic' (Akaike) or 'bic' (Bayesian).

    Returns
    -------
    n_terms : int
        Optimal number of Fourier terms.
    """
    y = np.asarray(y)
    n = len(y[~np.isnan(y)])

    if max_terms is None:
        max_terms = min(period // 2, 10)

    max_terms = max(1, min(max_terms, period // 2, n // 4))

    best_n = 1
    best_criterion = np.inf

    for n_terms in range(1, max_terms + 1):
        fourier = FourierSeasonality(period=period, n_terms=n_terms)

        try:
            fourier.fit(y)
        except ValueError:
            break

        t = np.arange(len(y))
        predicted = fourier.predict(t)

        # Calculate residual sum of squares
        mask = ~np.isnan(y)
        residuals = y[mask] - predicted[mask]
        rss = np.sum(residuals**2)

        # Number of parameters (2 per Fourier term)
        k = 2 * n_terms

        # Information criterion
        if criterion == "aic":
            # AIC = n * log(RSS/n) + 2k
            ic = n * np.log(rss / n) + 2 * k
        elif criterion == "bic":
            # BIC = n * log(RSS/n) + k * log(n)
            ic = n * np.log(rss / n) + k * np.log(n)
        else:
            raise ValueError(f"Unknown criterion: {criterion}")

        if ic < best_criterion:
            best_criterion = ic
            best_n = n_terms

    return best_n
