"""
换手率因子。
"""

import pandas as pd

from factors.base import BaseFactor


class TurnoverFactor(BaseFactor):
    """N 日均换手率。"""

    category = "technical"

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

    def calculate(self, df: pd.DataFrame) -> pd.Series:
        return df["turnover_rate"].rolling(window=self.period, min_periods=self.period).mean()

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


class TurnoverChangeFactor(BaseFactor):
    """换手率变化 = 当日换手率 / N 日均换手率。"""

    category = "technical"

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

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

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