"""Trend fitting algorithms for time series decomposition."""

from abc import ABC, abstractmethod
from dataclasses import dataclass
from typing import List, Literal, Optional, Tuple, Union

import numpy as np
import pandas as pd
from scipy import optimize
from statsmodels.nonparametric.smoothers_lowess import lowess


class BaseTrend(ABC):
    """Abstract base class for trend models."""

    @abstractmethod
    def fit(self, y: np.ndarray, t: np.ndarray) -> "BaseTrend":
        """Fit the trend model to data."""
        pass

    @abstractmethod
    def predict(self, t: Union[int, np.ndarray]) -> np.ndarray:
        """Predict trend values at given time indices."""
        pass

    @abstractmethod
    def get_fitted_values(self) -> np.ndarray:
        """Return fitted values for training period."""
        pass

    @property
    @abstractmethod
    def n_params(self) -> int:
        """Number of parameters in the model."""
        pass


@dataclass
class LinearTrend(BaseTrend):
    """
    Linear trend model: T(t) = alpha + beta * t

    Attributes
    ----------
    alpha : float
        Intercept parameter.
    beta : float
        Slope parameter.
    """

    alpha: float = 0.0
    beta: float = 0.0
    _fitted_values: Optional[np.ndarray] = None
    _t: Optional[np.ndarray] = None

    def fit(self, y: np.ndarray, t: Optional[np.ndarray] = None) -> "LinearTrend":
        """
        Fit linear trend using ordinary least squares.

        Parameters
        ----------
        y : np.ndarray
            Target values.
        t : np.ndarray, optional
            Time indices. Defaults to 0, 1, 2, ...

        Returns
        -------
        self : LinearTrend
            Fitted trend model.
        """
        n = len(y)
        if t is None:
            t = np.arange(n, dtype=float)

        self._t = t.copy()

        # Calculate means
        t_mean = np.mean(t)
        y_mean = np.mean(y)

        # Calculate slope (beta) using OLS
        numerator = np.sum((t - t_mean) * (y - y_mean))
        denominator = np.sum((t - t_mean) ** 2)

        if denominator == 0:
            self.beta = 0.0
        else:
            self.beta = numerator / denominator

        # Calculate intercept (alpha)
        self.alpha = y_mean - self.beta * t_mean

        # Store fitted values
        self._fitted_values = self.predict(t)

        return self

    def predict(self, t: Union[int, np.ndarray]) -> np.ndarray:
        """Predict trend values."""
        t = np.atleast_1d(t)
        return self.alpha + self.beta * t

    def get_fitted_values(self) -> np.ndarray:
        """Return fitted values."""
        if self._fitted_values is None:
            raise ValueError("Model must be fitted first")
        return self._fitted_values

    @property
    def n_params(self) -> int:
        """Number of parameters (alpha, beta)."""
        return 2


@dataclass
class ExponentialTrend(BaseTrend):
    """
    Exponential trend model: T(t) = alpha * exp(beta * t)

    Attributes
    ----------
    alpha : float
        Scale parameter.
    beta : float
        Growth rate parameter.
    """

    alpha: float = 1.0
    beta: float = 0.0
    _fitted_values: Optional[np.ndarray] = None
    _t: Optional[np.ndarray] = None

    def fit(self, y: np.ndarray, t: Optional[np.ndarray] = None) -> "ExponentialTrend":
        """
        Fit exponential trend using log-linear regression with bias correction.

        Parameters
        ----------
        y : np.ndarray
            Target values (must be positive).
        t : np.ndarray, optional
            Time indices. Defaults to 0, 1, 2, ...

        Returns
        -------
        self : ExponentialTrend
            Fitted trend model.

        Raises
        ------
        ValueError
            If y contains non-positive values.
        """
        n = len(y)
        if t is None:
            t = np.arange(n, dtype=float)

        self._t = t.copy()

        # Check for non-positive values
        if (y <= 0).any():
            raise ValueError("Exponential trend requires positive values")

        # Transform to log scale
        log_y = np.log(y)

        # Fit linear regression on log-transformed data
        linear = LinearTrend()
        linear.fit(log_y, t)

        log_alpha = linear.alpha
        self.beta = linear.beta

        # Bias correction for log transformation
        residuals = log_y - linear.predict(t)
        sigma_sq = np.var(residuals)
        self.alpha = np.exp(log_alpha + sigma_sq / 2)

        # Store fitted values
        self._fitted_values = self.predict(t)

        return self

    def predict(self, t: Union[int, np.ndarray]) -> np.ndarray:
        """Predict trend values."""
        t = np.atleast_1d(t)
        return self.alpha * np.exp(self.beta * t)

    def get_fitted_values(self) -> np.ndarray:
        """Return fitted values."""
        if self._fitted_values is None:
            raise ValueError("Model must be fitted first")
        return self._fitted_values

    @property
    def n_params(self) -> int:
        """Number of parameters (alpha, beta)."""
        return 2


@dataclass
class PolynomialTrend(BaseTrend):
    """
    Polynomial trend model: T(t) = sum(beta_j * t^j) for j = 0 to degree

    Attributes
    ----------
    coefficients : np.ndarray
        Polynomial coefficients [beta_0, beta_1, ..., beta_degree].
    degree : int
        Polynomial degree.
    """

    coefficients: Optional[np.ndarray] = None
    degree: int = 2
    _fitted_values: Optional[np.ndarray] = None
    _t: Optional[np.ndarray] = None

    def fit(self, y: np.ndarray, t: Optional[np.ndarray] = None) -> "PolynomialTrend":
        """
        Fit polynomial trend using ordinary least squares.

        Parameters
        ----------
        y : np.ndarray
            Target values.
        t : np.ndarray, optional
            Time indices. Defaults to 0, 1, 2, ...

        Returns
        -------
        self : PolynomialTrend
            Fitted trend model.
        """
        n = len(y)
        if t is None:
            t = np.arange(n, dtype=float)

        self._t = t.copy()

        # Normalize t to avoid numerical issues with high powers
        t_min, t_max = t.min(), t.max()
        if t_max > t_min:
            t_norm = (t - t_min) / (t_max - t_min)
        else:
            t_norm = t - t_min

        # Build design matrix
        X = np.column_stack([t_norm ** j for j in range(self.degree + 1)])

        # Solve normal equations using numpy's lstsq (more stable)
        self.coefficients, _, _, _ = np.linalg.lstsq(X, y, rcond=None)

        # Store normalization parameters for prediction
        self._t_min = t_min
        self._t_max = t_max

        # Store fitted values
        self._fitted_values = self.predict(t)

        return self

    def predict(self, t: Union[int, np.ndarray]) -> np.ndarray:
        """Predict trend values."""
        if self.coefficients is None:
            raise ValueError("Model must be fitted first")

        t = np.atleast_1d(t).astype(float)

        # Normalize t using same parameters
        if self._t_max > self._t_min:
            t_norm = (t - self._t_min) / (self._t_max - self._t_min)
        else:
            t_norm = t - self._t_min

        # Calculate polynomial
        result = np.zeros_like(t_norm)
        for j, coef in enumerate(self.coefficients):
            result += coef * (t_norm ** j)

        return result

    def get_fitted_values(self) -> np.ndarray:
        """Return fitted values."""
        if self._fitted_values is None:
            raise ValueError("Model must be fitted first")
        return self._fitted_values

    @property
    def n_params(self) -> int:
        """Number of parameters (degree + 1 coefficients)."""
        return self.degree + 1


@dataclass
class LOWESSTrend(BaseTrend):
    """
    LOWESS (Locally Weighted Scatterplot Smoothing) trend.

    A non-parametric method that fits local weighted regressions.

    Attributes
    ----------
    frac : float
        Fraction of data to use for each local regression (0-1).
    """

    frac: float = 0.3
    _fitted_values: Optional[np.ndarray] = None
    _t: Optional[np.ndarray] = None
    _y: Optional[np.ndarray] = None

    def fit(self, y: np.ndarray, t: Optional[np.ndarray] = None) -> "LOWESSTrend":
        """
        Fit LOWESS trend.

        Parameters
        ----------
        y : np.ndarray
            Target values.
        t : np.ndarray, optional
            Time indices. Defaults to 0, 1, 2, ...

        Returns
        -------
        self : LOWESSTrend
            Fitted trend model.
        """
        n = len(y)
        if t is None:
            t = np.arange(n, dtype=float)

        self._t = t.copy()
        self._y = y.copy()

        # Use statsmodels LOWESS
        smoothed = lowess(y, t, frac=self.frac, return_sorted=False)
        self._fitted_values = smoothed

        return self

    def predict(self, t: Union[int, np.ndarray]) -> np.ndarray:
        """
        Predict trend values.

        For LOWESS, we extrapolate using linear regression on the last
        portion of the fitted values for future predictions.
        """
        if self._fitted_values is None:
            raise ValueError("Model must be fitted first")

        t = np.atleast_1d(t).astype(float)
        result = np.zeros(len(t))

        t_max = self._t.max()
        t_min = self._t.min()

        for i, ti in enumerate(t):
            if ti <= t_max and ti >= t_min:
                # Interpolate within training range
                result[i] = np.interp(ti, self._t, self._fitted_values)
            elif ti > t_max:
                # Extrapolate forward using linear trend from last 10% of data
                n_tail = max(int(0.1 * len(self._t)), 2)
                tail_t = self._t[-n_tail:]
                tail_y = self._fitted_values[-n_tail:]

                # Fit linear trend to tail
                slope = (tail_y[-1] - tail_y[0]) / (tail_t[-1] - tail_t[0] + 1e-10)
                result[i] = tail_y[-1] + slope * (ti - t_max)
            else:
                # Extrapolate backward
                n_head = max(int(0.1 * len(self._t)), 2)
                head_t = self._t[:n_head]
                head_y = self._fitted_values[:n_head]

                slope = (head_y[-1] - head_y[0]) / (head_t[-1] - head_t[0] + 1e-10)
                result[i] = head_y[0] + slope * (ti - t_min)

        return result

    def get_fitted_values(self) -> np.ndarray:
        """Return fitted values."""
        if self._fitted_values is None:
            raise ValueError("Model must be fitted first")
        return self._fitted_values

    @property
    def n_params(self) -> int:
        """Number of parameters (approximate - based on effective df)."""
        # LOWESS doesn't have explicit parameters, return approx based on frac
        return max(2, int(self.frac * 10))


@dataclass
class DampedTrend(BaseTrend):
    """
    Damped linear trend: T(t+h) = L + sum(phi^i * b) for i=1 to h

    The trend converges asymptotically rather than growing indefinitely.

    Attributes
    ----------
    level : float
        Current level.
    trend : float
        Current trend component.
    phi : float
        Damping parameter (0 < phi < 1).
    """

    level: float = 0.0
    trend: float = 0.0
    phi: float = 0.98
    _fitted_values: Optional[np.ndarray] = None
    _t: Optional[np.ndarray] = None

    def fit(self, y: np.ndarray, t: Optional[np.ndarray] = None) -> "DampedTrend":
        """
        Fit damped trend using optimization.

        Parameters
        ----------
        y : np.ndarray
            Target values.
        t : np.ndarray, optional
            Time indices. Defaults to 0, 1, 2, ...

        Returns
        -------
        self : DampedTrend
            Fitted trend model.
        """
        n = len(y)
        if t is None:
            t = np.arange(n, dtype=float)

        self._t = t.copy()

        # Initial estimates from linear trend
        linear = LinearTrend()
        linear.fit(y, t)

        # Optimize phi, level, and trend
        def objective(params):
            level, trend, phi = params
            phi = np.clip(phi, 0.01, 0.99)

            fitted = np.zeros(n)
            fitted[0] = level

            for i in range(1, n):
                h = t[i] - t[0]
                phi_sum = np.sum([phi ** j for j in range(1, int(h) + 1)])
                fitted[i] = level + trend * phi_sum

            return np.sum((y - fitted) ** 2)

        # Initial guess
        x0 = [y[0], linear.beta, 0.98]
        bounds = [(None, None), (None, None), (0.01, 0.99)]

        result = optimize.minimize(objective, x0, bounds=bounds, method="L-BFGS-B")
        self.level, self.trend, self.phi = result.x

        # Store fitted values
        self._fitted_values = self.predict(t)

        return self

    def predict(self, t: Union[int, np.ndarray]) -> np.ndarray:
        """Predict trend values."""
        t = np.atleast_1d(t)
        result = np.zeros(len(t))

        t0 = self._t[0] if self._t is not None else 0

        for i, ti in enumerate(t):
            h = int(ti - t0)
            if h <= 0:
                result[i] = self.level
            else:
                phi_sum = np.sum([self.phi ** j for j in range(1, h + 1)])
                result[i] = self.level + self.trend * phi_sum

        return result

    def get_fitted_values(self) -> np.ndarray:
        """Return fitted values."""
        if self._fitted_values is None:
            raise ValueError("Model must be fitted first")
        return self._fitted_values

    @property
    def n_params(self) -> int:
        """Number of parameters (level, trend, phi)."""
        return 3


def select_best_trend(
    y: np.ndarray,
    t: Optional[np.ndarray] = None,
    candidate_types: List[str] = None,
    penalize_complexity: bool = True,
) -> BaseTrend:
    """
    Automatically select the best trend model.

    Parameters
    ----------
    y : np.ndarray
        Target values.
    t : np.ndarray, optional
        Time indices. Defaults to 0, 1, 2, ...
    candidate_types : List[str], optional
        List of trend types to try. Options: 'linear', 'exponential',
        'polynomial2', 'polynomial3', 'lowess', 'damped'.
        Defaults to ['linear', 'exponential', 'polynomial2'].
    penalize_complexity : bool, default True
        Whether to penalize complex models using AIC-like adjustment.

    Returns
    -------
    best_model : BaseTrend
        Best fitting trend model.

    Examples
    --------
    >>> y = np.array([100, 105, 112, 118, 125, 132])
    >>> trend = select_best_trend(y)
    >>> future = trend.predict(np.arange(6, 10))
    """
    n = len(y)
    if t is None:
        t = np.arange(n, dtype=float)

    if candidate_types is None:
        candidate_types = ["linear", "exponential", "polynomial2"]

    best_model = None
    best_score = np.inf

    for trend_type in candidate_types:
        try:
            if trend_type == "linear":
                model = LinearTrend()
            elif trend_type == "exponential":
                if (y <= 0).any():
                    continue  # Skip for non-positive data
                model = ExponentialTrend()
            elif trend_type == "polynomial2":
                model = PolynomialTrend(degree=2)
            elif trend_type == "polynomial3":
                model = PolynomialTrend(degree=3)
            elif trend_type == "lowess":
                model = LOWESSTrend(frac=0.3)
            elif trend_type == "damped":
                model = DampedTrend()
            else:
                continue

            model.fit(y, t)
            fitted = model.get_fitted_values()

            # Calculate SSE
            sse = np.sum((y - fitted) ** 2)

            # Penalize complexity (AIC-like adjustment)
            if penalize_complexity:
                k = model.n_params
                # Simple AIC approximation: n * log(SSE/n) + 2k
                if sse > 0:
                    score = n * np.log(sse / n) + 2 * k
                else:
                    score = -np.inf
            else:
                score = sse

            if score < best_score:
                best_score = score
                best_model = model

        except Exception:
            continue  # Skip models that fail to fit

    if best_model is None:
        # Fall back to linear trend
        best_model = LinearTrend()
        best_model.fit(y, t)

    return best_model


def fit_trend(
    y: Union[pd.Series, np.ndarray],
    trend_type: Literal["linear", "exponential", "polynomial", "lowess", "damped", "auto"] = "auto",
    degree: int = 2,
    frac: float = 0.3,
) -> Tuple[BaseTrend, np.ndarray]:
    """
    Convenience function to fit a trend to a time series.

    Parameters
    ----------
    y : pd.Series or np.ndarray
        Target time series.
    trend_type : str, default 'auto'
        Type of trend to fit.
    degree : int, default 2
        Polynomial degree (if trend_type='polynomial').
    frac : float, default 0.3
        LOWESS fraction (if trend_type='lowess').

    Returns
    -------
    trend : BaseTrend
        Fitted trend model.
    fitted_values : np.ndarray
        Fitted trend values for the training period.
    """
    y_arr = np.asarray(y)
    t = np.arange(len(y_arr), dtype=float)

    if trend_type == "auto":
        trend = select_best_trend(y_arr, t)
    elif trend_type == "linear":
        trend = LinearTrend()
        trend.fit(y_arr, t)
    elif trend_type == "exponential":
        trend = ExponentialTrend()
        trend.fit(y_arr, t)
    elif trend_type == "polynomial":
        trend = PolynomialTrend(degree=degree)
        trend.fit(y_arr, t)
    elif trend_type == "lowess":
        trend = LOWESSTrend(frac=frac)
        trend.fit(y_arr, t)
    elif trend_type == "damped":
        trend = DampedTrend()
        trend.fit(y_arr, t)
    else:
        raise ValueError(f"Unknown trend_type: {trend_type}")

    return trend, trend.get_fitted_values()
