"""
RSI 相对强弱因子。
"""

import pandas as pd

from factors.base import BaseFactor


class RSIFactor(BaseFactor):
    """Wilder's RSI = 100 - 100 / (1 + RS), RS = avg_gain / avg_loss"""

    category = "technical"

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

    def calculate(self, df: pd.DataFrame) -> pd.Series:
        delta = df["close"].diff()
        gain = delta.clip(lower=0)
        loss = (-delta).clip(lower=0)
        avg_gain = gain.ewm(span=self.period, min_periods=self.period).mean()
        avg_loss = loss.ewm(span=self.period, min_periods=self.period).mean()
        rs = avg_gain / avg_loss.replace(0, float("nan"))
        return 100 - 100 / (1 + rs)

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