"""
量价因子：量比 / 成交量变化率。
"""

import pandas as pd

from factors.base import BaseFactor


class VolumeFactor(BaseFactor):
    """N 日均量比 = vol / mean(vol, N)"""

    category = "technical"

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

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

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


class VolumeChangeFactor(BaseFactor):
    """成交量 N 日变化率"""

    category = "technical"

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

    def calculate(self, df: pd.DataFrame) -> pd.Series:
        return df["vol"].pct_change(periods=self.period) * 100

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