import pandas as pd

from signalab.indicator.mamethod import MAMethod


def _mid_price(candles: pd.DataFrame) -> pd.Series:
    """Compute candle mid-close price.

    Parameters
    ----------
    candles : pd.DataFrame
        Candle dataframe containing ``ask_close`` and ``bid_close``.

    Returns
    -------
    pd.Series
        Mid close series.
    """
    return (candles["ask_close"] + candles["bid_close"]) / 2


def _ema(series: pd.Series, period: int) -> pd.Series:
    """Compute exponential moving average.

    Parameters
    ----------
    series : pd.Series
        Input value series.
    period : int
        Smoothing period.

    Returns
    -------
    pd.Series
        EMA values.
    """
    return series.ewm(span=period, adjust=False).mean()


def _sma(series: pd.Series, period: int) -> pd.Series:
    """Compute simple moving average.

    Parameters
    ----------
    series : pd.Series
        Input value series.
    period : int
        Window length.

    Returns
    -------
    pd.Series
        SMA values.
    """
    return series.rolling(window=period).mean()


def _wma(series: pd.Series, period: int) -> pd.Series:
    """Compute linearly weighted moving average.

    Parameters
    ----------
    series : pd.Series
        Input value series.
    period : int
        Window length.

    Returns
    -------
    pd.Series
        WMA values.
    """
    weights = pd.Series(range(1, period + 1))
    return series.rolling(period).apply(
        lambda x: (x * weights).sum() / weights.sum(), raw=True
    )


def _smma(series: pd.Series, period: int) -> pd.Series:
    """Compute smoothed moving average.

    Parameters
    ----------
    series : pd.Series
        Input value series.
    period : int
        Smoothing period.

    Returns
    -------
    pd.Series
        SMMA values, equivalent to EMA with ``alpha=1/period``.
    """
    return series.ewm(alpha=1.0 / period, adjust=False).mean()


def _ma(series: pd.Series, period: int, method: MAMethod) -> pd.Series:
    """Dispatch moving-average computation by method.

    Parameters
    ----------
    series : pd.Series
        Input value series.
    period : int
        Window or smoothing period.
    method : MAMethod
        Moving-average implementation.

    Returns
    -------
    pd.Series
        Moving-average values.

    Raises
    ------
    ValueError
        If ``method`` is not supported.
    """
    if method == MAMethod.SMA:
        return _sma(series, period)
    elif method == MAMethod.EMA:
        return _ema(series, period)
    elif method == MAMethod.WMA:
        return _wma(series, period)
    elif method == MAMethod.SMMA:
        return _smma(series, period)
    raise ValueError(f"Unknown MA method: {method}")
