"""
波动率因子。
"""

import pandas as pd

from factors.base import BaseFactor


class VolatilityFactor(BaseFactor):
    """N 日年化波动率 = std(daily_return, N) * sqrt(252) * 100"""

    category = "technical"

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

    def calculate(self, df: pd.DataFrame) -> pd.Series:
        daily_ret = df["close"].pct_change()
        return daily_ret.rolling(window=self.period, min_periods=self.period).std() * (252 ** 0.5) * 100

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


class DownsideVolatilityFactor(BaseFactor):
    """下行波动率：只计算负收益的标准差。"""

    category = "technical"

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

    def calculate(self, df: pd.DataFrame) -> pd.Series:
        daily_ret = df["close"].pct_change()
        downside = daily_ret.clip(upper=0)
        return downside.rolling(window=self.period, min_periods=self.period).std() * (252 ** 0.5) * 100

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