"""
MACD 因子。
"""

import pandas as pd

from factors.base import BaseFactor


class MACDFactor(BaseFactor):
    """
    MACD 系列因子。

    返回 DIF/DEA/HIST 三个值。使用 calculate() 返回 HIST（柱），
    单独方法获取 DIF/DEA。
    """

    category = "technical"

    def __init__(self, fast: int = 12, slow: int = 26, signal: int = 9):
        self.fast = fast
        self.slow = slow
        self.signal = signal
        self.name = f"macd_{fast}_{slow}_{signal}"

    def _ema(self, series: pd.Series, span: int) -> pd.Series:
        return series.ewm(span=span, min_periods=span).mean()

    def calculate(self, df: pd.DataFrame) -> pd.Series:
        """返回 MACD 柱（DIF - DEA）。"""
        ema_fast = self._ema(df["close"], self.fast)
        ema_slow = self._ema(df["close"], self.slow)
        dif = ema_fast - ema_slow
        dea = self._ema(dif, self.signal)
        return dif - dea

    def dif(self, df: pd.DataFrame) -> pd.Series:
        ema_fast = self._ema(df["close"], self.fast)
        ema_slow = self._ema(df["close"], self.slow)
        return ema_fast - ema_slow

    def dea(self, df: pd.DataFrame) -> pd.Series:
        dif = self.dif(df)
        return self._ema(dif, self.signal)

    def get_required_columns(self) -> list[str]:
        return ["close"]