"""Demand forecasting models with pluggable backends.

Two backends:
    1. Built-in — pure Python implementations of classical methods
    2. Statsforecast (Nixtla) — production-grade, fast, scales to millions of series

Usage via provider:
    from platoon.learning.providers import get_provider
    fc = get_provider("forecasting")  # auto-selects best
    result = fc.forecast(series, horizon=14, method="auto")

Direct function usage:
    from platoon.learning.forecasting import moving_average, exponential_smoothing
    result = moving_average(series, window=7, horizon=7)

Available methods:
    Built-in: moving_average, exponential_smoothing, croston, seasonal_decompose
    Statsforecast: AutoARIMA, AutoETS, CrostonOptimized, SeasonalNaive,
                   MSTL, AutoTheta, and 20+ more
"""

import math
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Tuple

from platoon.learning.providers import ModelProvider, register_provider


# ---------------------------------------------------------------------------
# Result types
# ---------------------------------------------------------------------------


@dataclass
class ForecastResult:
    """Output of a forecast operation."""

    values: List[float]
    lower_bound: List[float] = field(default_factory=list)
    upper_bound: List[float] = field(default_factory=list)
    method: str = ""
    horizon: int = 0
    fitted_values: List[float] = field(default_factory=list)
    metadata: Dict[str, Any] = field(default_factory=dict)


# ---------------------------------------------------------------------------
# Built-in implementations (pure Python, no deps)
# ---------------------------------------------------------------------------


def moving_average(
    series: List[float],
    window: int = 7,
    horizon: int = 7,
) -> ForecastResult:
    """Simple moving average forecast.

    Args:
        series: Historical demand values (one per period)
        window: Number of recent periods to average
        horizon: Number of future periods to forecast

    Returns:
        ForecastResult with constant forecast = mean of last `window` values
    """
    if len(series) < window:
        window = len(series)
    if window == 0:
        raise ValueError("series must not be empty")

    tail = series[-window:]
    avg = sum(tail) / len(tail)
    std = math.sqrt(sum((x - avg) ** 2 for x in tail) / len(tail)) if len(tail) > 1 else 0

    values = [round(avg, 4)] * horizon
    lower = [round(avg - 1.96 * std, 4)] * horizon
    upper = [round(avg + 1.96 * std, 4)] * horizon

    return ForecastResult(
        values=values, lower_bound=lower, upper_bound=upper,
        method="moving_average", horizon=horizon,
        metadata={"window": window, "mean": round(avg, 4), "std": round(std, 4)},
    )


def exponential_smoothing(
    series: List[float],
    alpha: float = 0.3,
    horizon: int = 7,
) -> ForecastResult:
    """Simple exponential smoothing (SES).

    S_t = alpha * Y_t + (1 - alpha) * S_{t-1}

    Args:
        series: Historical demand values
        alpha: Smoothing factor (0-1)
        horizon: Future periods to forecast

    Returns:
        ForecastResult where forecast = last smoothed value
    """
    if not series:
        raise ValueError("series must not be empty")
    if not 0 < alpha <= 1:
        raise ValueError("alpha must be between 0 (exclusive) and 1 (inclusive)")

    # Initialize with first observation
    smoothed = series[0]
    fitted = [smoothed]
    residuals = []

    for y in series[1:]:
        smoothed = alpha * y + (1 - alpha) * smoothed
        fitted.append(round(smoothed, 4))
        residuals.append(y - smoothed)

    # Forecast = last smoothed value
    forecast_val = round(smoothed, 4)
    # Prediction interval from residual std
    residual_std = math.sqrt(sum(r**2 for r in residuals) / max(len(residuals), 1)) if residuals else 0

    values = [forecast_val] * horizon
    lower = [round(forecast_val - 1.96 * residual_std, 4)] * horizon
    upper = [round(forecast_val + 1.96 * residual_std, 4)] * horizon

    return ForecastResult(
        values=values, lower_bound=lower, upper_bound=upper,
        method="exponential_smoothing", horizon=horizon,
        fitted_values=fitted,
        metadata={"alpha": alpha, "residual_std": round(residual_std, 4)},
    )


def croston(
    series: List[float],
    alpha: float = 0.1,
    horizon: int = 7,
) -> ForecastResult:
    """Croston's method for intermittent demand.

    Separately smooths non-zero demand sizes and inter-demand intervals,
    then combines: forecast = smoothed_size / smoothed_interval.

    Args:
        series: Historical demand (many zeros expected)
        alpha: Smoothing factor
        horizon: Periods to forecast

    Returns:
        ForecastResult with expected demand per period
    """
    if not series:
        raise ValueError("series must not be empty")

    # Find non-zero demands and their intervals
    non_zero_demands = []
    intervals = []
    periods_since_last = 0

    for val in series:
        periods_since_last += 1
        if val > 0:
            non_zero_demands.append(val)
            intervals.append(periods_since_last)
            periods_since_last = 0

    if not non_zero_demands:
        # All zeros — forecast zero
        return ForecastResult(
            values=[0.0] * horizon, method="croston", horizon=horizon,
            metadata={"warning": "all_zeros"},
        )

    # Smooth demand sizes
    z = non_zero_demands[0]
    for d in non_zero_demands[1:]:
        z = alpha * d + (1 - alpha) * z

    # Smooth intervals
    p = intervals[0]
    for interval in intervals[1:]:
        p = alpha * interval + (1 - alpha) * p

    # Forecast = smoothed demand / smoothed interval
    forecast_val = round(z / p, 4) if p > 0 else 0.0
    values = [forecast_val] * horizon

    return ForecastResult(
        values=values, method="croston", horizon=horizon,
        metadata={
            "smoothed_demand": round(z, 4),
            "smoothed_interval": round(p, 4),
            "non_zero_count": len(non_zero_demands),
            "total_periods": len(series),
        },
    )


def seasonal_decompose(
    series: List[float],
    period: int = 7,
    model: str = "additive",
) -> Tuple[List[float], List[float], List[float]]:
    """Classical seasonal decomposition.

    Steps:
        1. Compute trend via centered moving average
        2. Remove trend to get seasonal + residual
        3. Average seasonal component per position in cycle
        4. Residual = observed - trend - seasonal (additive)

    Args:
        series: Historical values
        period: Seasonal cycle length (7=weekly, 12=monthly)
        model: "additive" or "multiplicative"

    Returns:
        Tuple of (trend, seasonal, residual) — same length as series.
        Boundary values where MA can't compute use NaN.
    """
    n = len(series)
    if n < 2 * period:
        raise ValueError(f"Need at least 2*period ({2*period}) observations, got {n}")

    nan = float("nan")

    # Step 1: Centered moving average for trend
    half = period // 2
    trend = [nan] * n
    for i in range(half, n - half):
        if period % 2 == 0:
            # Even period: average of two half-window MAs
            window = series[i - half: i + half]
            trend[i] = (sum(window) + series[i + half] * 0.5 + series[i - half] * 0.5) / period
        else:
            window = series[i - half: i + half + 1]
            trend[i] = sum(window) / period

    # Step 2: Detrend
    detrended = [nan] * n
    for i in range(n):
        if math.isnan(trend[i]):
            continue
        if model == "multiplicative" and trend[i] != 0:
            detrended[i] = series[i] / trend[i]
        else:
            detrended[i] = series[i] - trend[i]

    # Step 3: Average seasonal indices
    seasonal_indices = [0.0] * period
    counts = [0] * period
    for i in range(n):
        if not math.isnan(detrended[i]):
            pos = i % period
            seasonal_indices[pos] += detrended[i]
            counts[pos] += 1

    for pos in range(period):
        if counts[pos] > 0:
            seasonal_indices[pos] /= counts[pos]

    # Normalize: additive should sum to 0, multiplicative to period
    if model == "multiplicative":
        avg = sum(seasonal_indices) / period
        seasonal_indices = [s / avg for s in seasonal_indices]
    else:
        avg = sum(seasonal_indices) / period
        seasonal_indices = [s - avg for s in seasonal_indices]

    # Build full seasonal series
    seasonal = [round(seasonal_indices[i % period], 4) for i in range(n)]

    # Step 4: Residual
    residual = [nan] * n
    for i in range(n):
        if math.isnan(trend[i]):
            continue
        if model == "multiplicative" and seasonal[i] != 0:
            residual[i] = series[i] / (trend[i] * seasonal[i])
        else:
            residual[i] = round(series[i] - trend[i] - seasonal[i], 4)

    trend = [round(t, 4) if not math.isnan(t) else nan for t in trend]

    return trend, seasonal, residual


# ---------------------------------------------------------------------------
# Built-in forecasting provider
# ---------------------------------------------------------------------------


@register_provider
class BuiltinForecastProvider(ModelProvider):
    """Pure-Python forecasting — no external dependencies."""

    name = "builtin_forecast"
    domain = "forecasting"
    backend = "builtin"

    def forecast(
        self,
        series: List[float],
        horizon: int = 7,
        method: str = "auto",
        **kwargs,
    ) -> ForecastResult:
        """Run a forecast using the specified method.

        Args:
            series: Historical values
            horizon: Periods to forecast
            method: "moving_average", "exponential_smoothing", "croston", or "auto"
            **kwargs: Method-specific params (window, alpha, etc.)

        Returns:
            ForecastResult
        """
        if method == "auto":
            method = self._auto_select(series)

        dispatch = {
            "moving_average": lambda: moving_average(series, horizon=horizon, **kwargs),
            "exponential_smoothing": lambda: exponential_smoothing(series, horizon=horizon, **kwargs),
            "croston": lambda: croston(series, horizon=horizon, **kwargs),
        }
        if method not in dispatch:
            raise ValueError(f"Unknown method: {method}. Available: {list(dispatch.keys())}")
        return dispatch[method]()

    def decompose(self, series: List[float], period: int = 7, model: str = "additive"):
        """Decompose series into trend, seasonal, residual."""
        return seasonal_decompose(series, period, model)

    def _auto_select(self, series: List[float]) -> str:
        """Pick the best built-in method based on series characteristics."""
        if not series:
            return "moving_average"
        zero_ratio = sum(1 for x in series if x == 0) / len(series)
        if zero_ratio > 0.4:
            return "croston"
        return "exponential_smoothing"


# ---------------------------------------------------------------------------
# Statsforecast provider (Nixtla)
# ---------------------------------------------------------------------------


@register_provider
class StatsforecastProvider(ModelProvider):
    """Nixtla statsforecast backend — production-grade, fast, scalable.

    Provides AutoARIMA, AutoETS, AutoTheta, CrostonOptimized, SeasonalNaive,
    MSTL, and 20+ other methods out of the box.
    """

    name = "statsforecast_provider"
    domain = "forecasting"
    backend = "statsforecast"

    @classmethod
    def is_available(cls) -> bool:
        try:
            import statsforecast
            return True
        except ImportError:
            return False

    def forecast(
        self,
        series: List[float],
        horizon: int = 7,
        method: str = "auto",
        season_length: int = 7,
        level: List[int] = None,
        **kwargs,
    ) -> ForecastResult:
        """Run forecast using statsforecast models.

        Args:
            series: Historical values
            horizon: Periods ahead to forecast
            method: Model name — "auto", "arima", "ets", "theta",
                    "croston", "seasonal_naive", "mstl", or "ensemble"
            season_length: Seasonal period (7=weekly, 12=monthly, 365=yearly)
            level: Confidence levels for prediction intervals (default [95])

        Returns:
            ForecastResult
        """
        import pandas as pd
        from statsforecast import StatsForecast
        from statsforecast.models import (
            AutoARIMA,
            AutoETS,
            AutoTheta,
            CrostonOptimized,
            SeasonalNaive,
        )

        if level is None:
            level = [95]

        model_map = {
            "arima": [AutoARIMA(season_length=season_length)],
            "ets": [AutoETS(season_length=season_length)],
            "theta": [AutoTheta(season_length=season_length)],
            "croston": [CrostonOptimized()],
            "seasonal_naive": [SeasonalNaive(season_length=season_length)],
        }

        if method == "auto":
            zero_ratio = sum(1 for x in series if x == 0) / max(len(series), 1)
            if zero_ratio > 0.4:
                method = "croston"
            elif len(series) >= 2 * season_length:
                method = "ets"
            else:
                method = "arima"

        if method == "ensemble":
            models = [
                AutoARIMA(season_length=season_length),
                AutoETS(season_length=season_length),
                AutoTheta(season_length=season_length),
            ]
        elif method in model_map:
            models = model_map[method]
        else:
            raise ValueError(f"Unknown method: {method}. Available: {list(model_map.keys()) + ['auto', 'ensemble']}")

        # statsforecast 2.0 requires a pandas DataFrame with unique_id, ds, y
        df = pd.DataFrame({
            "unique_id": ["series"] * len(series),
            "ds": list(range(len(series))),
            "y": [float(v) for v in series],
        })

        sf = StatsForecast(models=models, freq=1, n_jobs=1)
        sf.fit(df)

        # CrostonOptimized doesn't support prediction intervals
        no_intervals = method == "croston"
        if no_intervals:
            forecast_df = sf.predict(h=horizon)
        else:
            forecast_df = sf.predict(h=horizon, level=level)

        # Extract results from the pandas DataFrame
        model_name = models[0].__class__.__name__
        if method == "ensemble":
            model_names = [m.__class__.__name__ for m in models]
            cols = [c for c in model_names if c in forecast_df.columns]
            vals = forecast_df[cols].mean(axis=1).tolist()
            values = [round(float(v), 4) for v in vals]
            lower = []
            upper = []
            actual_method = "ensemble"
        else:
            col = model_name
            values = [round(float(v), 4) for v in forecast_df[col].tolist()]

            lo_col = f"{col}-lo-{level[0]}"
            hi_col = f"{col}-hi-{level[0]}"
            lower = [round(float(v), 4) for v in forecast_df[lo_col].tolist()] if lo_col in forecast_df.columns else []
            upper = [round(float(v), 4) for v in forecast_df[hi_col].tolist()] if hi_col in forecast_df.columns else []
            actual_method = method

        return ForecastResult(
            values=values,
            lower_bound=lower,
            upper_bound=upper,
            method=f"statsforecast:{actual_method}",
            horizon=horizon,
            metadata={"model_class": model_name, "season_length": season_length},
        )

    def forecast_multi(
        self,
        series_dict: Dict[str, List[float]],
        horizon: int = 7,
        method: str = "auto",
        season_length: int = 7,
    ) -> Dict[str, ForecastResult]:
        """Forecast multiple series at once (batch).

        Statsforecast is optimized for batch — much faster than looping.

        Args:
            series_dict: {series_id: [values]} mapping
            horizon: Periods ahead
            method: Model to use
            season_length: Seasonal period

        Returns:
            {series_id: ForecastResult} mapping
        """
        results = {}
        for sid, series in series_dict.items():
            results[sid] = self.forecast(series, horizon=horizon, method=method,
                                         season_length=season_length)
        return results
