"""
均线交叉因子。
"""

import pandas as pd

from factors.base import BaseFactor


class MACrossFactor(BaseFactor):
    """
    均线交叉信号。

    返回：fast_ma / slow_ma - 1，正值表示短期均线在上方。
    """

    category = "technical"

    def __init__(self, fast: int = 5, slow: int = 20):
        self.fast = fast
        self.slow = slow
        self.name = f"ma_cross_{fast}_{slow}"

    def calculate(self, df: pd.DataFrame) -> pd.Series:
        ma_fast = df["close"].rolling(window=self.fast, min_periods=self.fast).mean()
        ma_slow = df["close"].rolling(window=self.slow, min_periods=self.slow).mean()
        return (ma_fast / ma_slow.replace(0, float("nan"))) - 1

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


class MADeviationFactor(BaseFactor):
    """
    价格偏离均线程度 = (close - ma) / ma * 100
    """

    category = "technical"

    def __init__(self, period: int = 20):
        self.period = period
        self.name = f"ma_dev_{period}"

    def calculate(self, df: pd.DataFrame) -> pd.Series:
        ma = df["close"].rolling(window=self.period, min_periods=self.period).mean()
        return (df["close"] - ma) / ma.replace(0, float("nan")) * 100

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