"""
振幅因子。
"""

import pandas as pd

from factors.base import BaseFactor


class AmplitudeFactor(BaseFactor):
    """N 日均振幅 = mean((high - low) / close, N) * 100"""

    category = "technical"

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

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

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