import pandas as pd

from signalab.indicator.mamethod import MAMethod


def _mid_price(candles: pd.DataFrame) -> pd.Series:
    return (candles["ask_close"] + candles["bid_close"]) / 2


def _ema(series: pd.Series, period: int) -> pd.Series:
    return series.ewm(span=period, adjust=False).mean()


def _sma(series: pd.Series, period: int) -> pd.Series:
    return series.rolling(window=period).mean()


def _wma(series: pd.Series, period: int) -> pd.Series:
    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:
    """Smoothed Moving Average — 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:
    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}")
