"""
布林带因子：价格在布林带中的位置。
"""

import pandas as pd

from factors.base import BaseFactor


class BollingerFactor(BaseFactor):
    """
    布林带位置 = (close - middle) / (upper - lower)

    值在 0~1 之间：接近 0 表示在下轨，接近 1 表示在上轨。
    """

    category = "technical"

    def __init__(self, period: int = 20, std: float = 2.0):
        self.period = period
        self.std = std
        self.name = f"boll_{period}"

    def calculate(self, df: pd.DataFrame) -> pd.Series:
        middle = df["close"].rolling(window=self.period, min_periods=self.period).mean()
        std = df["close"].rolling(window=self.period, min_periods=self.period).std()
        upper = middle + self.std * std
        lower = middle - self.std * std
        band_width = upper - lower
        return ((df["close"] - lower) / band_width.replace(0, float("nan"))).clip(0, 1)

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


class BollingerWidthFactor(BaseFactor):
    """布林带宽度 = (upper - lower) / middle * 100"""

    category = "technical"

    def __init__(self, period: int = 20, std: float = 2.0):
        self.period = period
        self.std = std
        self.name = f"boll_width_{period}"

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

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