"""
Time Series — Trend decomposition, forecasting, and seasonality detection.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Literal

import numpy as np
import pandas as pd
from statsmodels.tsa.holtwinters import ExponentialSmoothing
from statsmodels.tsa.seasonal import seasonal_decompose
from statsmodels.tsa.statespace.sarimax import SARIMAX
from sklearn.linear_model import LinearRegression
from sklearn.metrics import mean_absolute_error, mean_absolute_percentage_error

from bizanalyticsai.core.dataframe import BizFrame


@dataclass
class ForecastResult:
    """Result of a time-series forecast."""

    forecast:     BizFrame      # columns: date, forecast, lower_ci, upper_ci
    model_name:   str
    mae:          float | None  # in-sample MAE
    mape:         float | None  # in-sample MAPE (%)
    horizon:      int
    params:       dict = field(default_factory=dict)

    def summary(self) -> str:
        lines = [
            f"Model: {self.model_name} | Horizon: {self.horizon} periods",
        ]
        if self.mae is not None:
            lines.append(f"MAE: {self.mae:.4f}  |  MAPE: {self.mape:.2f}%")
        lines.append(f"\n{self.forecast.to_string(index=False)}")
        return "\n".join(lines)


def trend(
    series: pd.Series,
    method: Literal["linear", "polynomial", "ma"] = "linear",
    degree: int = 2,
    window: int = 7,
) -> BizFrame:
    """
    Fit a trend line to a time series.

    Args:
        series: Numeric time series.
        method: 'linear', 'polynomial', or 'ma' (moving average).
        degree: Polynomial degree (only for 'polynomial').
        window: Window size for moving average.

    Returns:
        BizFrame with original, trend, and residual columns.
    """
    s = series.reset_index(drop=True)
    x = np.arange(len(s)).reshape(-1, 1)

    if method == "ma":
        trend_vals = s.rolling(window=window, center=True, min_periods=1).mean()
    elif method == "linear":
        model = LinearRegression().fit(x, s)
        trend_vals = pd.Series(model.predict(x), index=s.index)
    else:  # polynomial
        coeffs     = np.polyfit(x.flatten(), s, degree)
        trend_vals = pd.Series(np.polyval(coeffs, x.flatten()), index=s.index)

    return BizFrame({
        "value":    s.values,
        "trend":    trend_vals.values,
        "residual": (s - trend_vals).values,
    })


def decompose(
    series: pd.Series,
    period: int = 12,
    model: Literal["additive", "multiplicative"] = "additive",
) -> dict:
    """
    Seasonal decomposition of a time series (STL-based).

    Returns dict with trend, seasonal, and residual Series.
    """
    result = seasonal_decompose(series.dropna(), model=model, period=period, extrapolate_trend="freq")
    return {
        "trend":    result.trend,
        "seasonal": result.seasonal,
        "residual": result.resid,
        "observed": result.observed,
    }


def forecast(
    series: pd.Series,
    horizon: int = 12,
    method: Literal["holtwinters", "arima", "sarima", "linear", "naive"] = "holtwinters",
    confidence: float = 0.95,
    **kwargs,
) -> ForecastResult:
    """
    Forecast future values of a time series.

    Methods:
    - holtwinters: Holt-Winters triple exponential smoothing (trend + seasonality)
    - arima:       ARIMA(p,d,q) — auto-selected if not provided
    - sarima:      Seasonal ARIMA (SARIMAX) with configurable PDQ
    - linear:      Simple linear extrapolation
    - naive:       Last value repeated (baseline)

    Args:
        series:    Historical time series (datetime index recommended).
        horizon:   Number of periods to forecast.
        method:    Forecasting algorithm.
        confidence: Confidence interval level (0.80–0.99).
        **kwargs:  Additional parameters for the underlying model.

    Returns:
        ForecastResult with forecast DataFrame and evaluation metrics.
    """
    s     = series.dropna()
    alpha = 1 - confidence
    preds: pd.Series | np.ndarray
    lower = upper = None
    params: dict = {}
    mae = mape = None

    # Build future index
    if isinstance(s.index, pd.DatetimeIndex):
        freq  = pd.infer_freq(s.index) or "ME"
        future_idx = pd.date_range(s.index[-1], periods=horizon + 1, freq=freq)[1:]
    else:
        future_idx = pd.RangeIndex(start=len(s), stop=len(s) + horizon)

    if method == "naive":
        last  = float(s.iloc[-1])
        preds = np.full(horizon, last)

    elif method == "linear":
        x     = np.arange(len(s)).reshape(-1, 1)
        model = LinearRegression().fit(x, s.values)
        x_fut = np.arange(len(s), len(s) + horizon).reshape(-1, 1)
        preds = model.predict(x_fut)
        mae   = float(mean_absolute_error(s.values, model.predict(x)))
        mape  = float(mean_absolute_percentage_error(s.values, model.predict(x))) * 100
        params = {"coef": float(model.coef_[0]), "intercept": float(model.intercept_)}

    elif method == "holtwinters":
        seasonal_periods = kwargs.pop("seasonal_periods", max(4, len(s) // 4))
        trend_hw   = kwargs.pop("trend", "add")
        seasonal   = kwargs.pop("seasonal", "add" if len(s) >= seasonal_periods * 2 else None)
        model = ExponentialSmoothing(
            s, trend=trend_hw, seasonal=seasonal,
            seasonal_periods=seasonal_periods if seasonal else None,
            **kwargs,
        ).fit(optimized=True)
        fc_obj = model.forecast(horizon)
        preds  = fc_obj.values
        # Confidence intervals (simulated via residual std)
        resid_std = float(model.resid.std())
        z = 1.96 if confidence >= 0.95 else 1.645
        spread = np.arange(1, horizon + 1) ** 0.5 * resid_std * z
        lower  = preds - spread
        upper  = preds + spread
        mae  = float(mean_absolute_error(s.values, model.fittedvalues))
        mape = float(mean_absolute_percentage_error(s.values, model.fittedvalues)) * 100

    elif method in ("arima", "sarima"):
        order   = kwargs.pop("order", (1, 1, 1))
        s_order = kwargs.pop("seasonal_order", (1, 1, 0, 12) if method == "sarima" else (0, 0, 0, 0))
        model   = SARIMAX(s, order=order, seasonal_order=s_order, **kwargs).fit(disp=False)
        fc_obj  = model.get_forecast(steps=horizon)
        preds   = fc_obj.predicted_mean.values
        ci      = fc_obj.conf_int(alpha=alpha)
        lower   = ci.iloc[:, 0].values
        upper   = ci.iloc[:, 1].values
        fitted  = model.fittedvalues
        mae  = float(mean_absolute_error(s.values, fitted))
        mape = float(mean_absolute_percentage_error(s.values, fitted)) * 100
        params = {"order": order, "seasonal_order": s_order}

    else:
        raise ValueError(f"Unknown method: {method}")

    fc_df = BizFrame({
        "date":     list(future_idx),
        "forecast": np.round(preds, 4).tolist(),
        "lower_ci": (np.round(lower, 4).tolist() if lower is not None else [None] * horizon),
        "upper_ci": (np.round(upper, 4).tolist() if upper is not None else [None] * horizon),
    })

    return ForecastResult(
        forecast=fc_df,
        model_name=method,
        mae=round(mae, 4) if mae is not None else None,
        mape=round(mape, 4) if mape is not None else None,
        horizon=horizon,
        params=params,
    )


def seasonality(series: pd.Series, period: int = 12) -> dict:
    """
    Test for seasonality and return seasonal indices.

    Returns dict with: has_seasonality, seasonal_indices, dominant_period.
    """
    s = series.dropna()
    if len(s) < period * 2:
        return {"has_seasonality": False, "seasonal_indices": [], "dominant_period": None}

    decomp = seasonal_decompose(s, model="additive", period=period, extrapolate_trend="freq")
    seasonal = decomp.seasonal.dropna()
    # Compute one full cycle of seasonal indices
    indices = seasonal.values[:period].tolist()
    strength = float(np.std(indices) / (np.std(s) + 1e-9))

    return {
        "has_seasonality":  strength > 0.1,
        "seasonal_strength": round(strength, 4),
        "seasonal_indices":  [round(v, 4) for v in indices],
        "dominant_period":   period,
    }


def growth_rate(series: pd.Series, periods: int = 1) -> BizFrame:
    """
    Compute absolute and percentage period-over-period growth.

    Returns BizFrame with original value, abs_change, pct_change, cagr.
    """
    s = series.reset_index(drop=True)
    n = len(s)
    cagr = float(((s.iloc[-1] / max(s.iloc[0], 1e-9)) ** (1 / max(n - 1, 1)) - 1) * 100)

    return BizFrame({
        "value":      s.values,
        "abs_change": s.diff(periods).values,
        "pct_change": (s.pct_change(periods) * 100).values,
        "cagr_pct":   [cagr] * n,
    })
