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

import numpy as np
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()
        # 标准 Wilder RSI：avg_loss==0 时 RSI 应 = 100，而非 NaN。
        # 用 where 显式处理除零，避免 replace(0, nan) 把上涨趋势判为缺失。
        rs = avg_gain / avg_loss.where(avg_loss != 0, np.nan)
        rsi = 100 - 100 / (1 + rs)
        # 上涨且无下跌的高位情形补 100（无 prior-loss 的窗口仍留 NaN 由上游填充）
        rsi = rsi.where(avg_loss != 0, 100.0)
        return rsi

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