"""
特殊因子模块
包含基于股票截面的特殊因子实现
"""

import numpy as np
import polars as pl
from main.factor.operator_framework import StockWiseFactor


class LimitFactor(StockWiseFactor):
    """涨跌停因子"""
    
    def __init__(self):
        super().__init__(
            name="limit",
            parameters={},
            required_factor_ids=["close", "up_limit", "down_limit"]
        )

    def calc_factor(self, group_df: pl.DataFrame) -> pl.Series:
        # 计算是否涨停或跌停
        close = group_df["close"]
        up_limit = group_df["up_limit"]
        down_limit = group_df["down_limit"]
        
        # 是否涨停
        is_up_limit = (close == up_limit).cast(pl.Int32)
        # 是否跌停
        is_down_limit = (close == down_limit).cast(pl.Int32)
        
        # 合并为一个因子
        limit_factor = is_up_limit - is_down_limit
        return limit_factor.alias(self.factor_id)


class VolumeRatioFactor(StockWiseFactor):
    """量比因子"""
    
    def __init__(self):
        super().__init__(
            name="volume_ratio",
            parameters={},
            required_factor_ids=["vol"]
        )

    def calc_factor(self, group_df: pl.DataFrame) -> pl.Series:
        # 计算量比：当日成交量 / 5日平均成交量
        vol = group_df["vol"]
        avg_vol_5d = vol.rolling_mean(5)
        # 避免除零
        volume_ratio = vol / (avg_vol_5d + 1e-8)
        return volume_ratio.alias(self.factor_id)


class BBI_RATIO_FACTOR(StockWiseFactor):
    """BBI比率因子"""
    
    def __init__(self):
        super().__init__(
            name="bbi_ratio",
            parameters={},
            required_factor_ids=["close"]
        )

    def calc_factor(self, group_df: pl.DataFrame) -> pl.Series:
        # 计算BBI比率
        close = group_df["close"]
        
        # 计算不同周期的SMA
        sma3 = close.rolling_mean(3)
        sma6 = close.rolling_mean(6)
        sma12 = close.rolling_mean(12)
        sma24 = close.rolling_mean(24)
        
        # 计算BBI
        bbi = (sma3 + sma6 + sma12 + sma24) / 4
        
        # 计算BBI比率
        bbi_ratio = bbi / (close + 1e-8)  # 避免除零
        
        return bbi_ratio.alias(self.factor_id)


class VolatilitySlopeFactor(StockWiseFactor):
    """波动率斜率因子"""
    
    def __init__(self, window_vol: int = 20, window_slope: int = 5):
        super().__init__(
            name="volatility_slope",
            parameters={"window_vol": window_vol, "window_slope": window_slope},
            required_factor_ids=["pct_chg"]
        )

    def calc_factor(self, group_df: pl.DataFrame) -> pl.Series:
        window_vol = self.parameters["window_vol"]
        window_slope = self.parameters["window_slope"]
        
        # 计算滚动标准差
        volatility = group_df["pct_chg"].rolling_std(window_vol)
        
        # 计算斜率
        # 这里简化处理，直接计算最后一个窗口的斜率
        # 实际应用中可能需要更复杂的线性回归计算
        volatility_slope = volatility.diff().alias(self.factor_id)
        
        return volatility_slope


class PriceVolumeTrendFactor(StockWiseFactor):
    """价格量能趋势因子"""
    
    def __init__(self):
        super().__init__(
            name="price_volume_trend",
            parameters={},
            required_factor_ids=["close", "vol"]
        )

    def calc_factor(self, group_df: pl.DataFrame) -> pl.Series:
        # 计算价格量能趋势
        close = group_df["close"]
        vol = group_df["vol"]
        
        # 计算价格变化
        price_change = close.diff()
        
        # 计算成交量变化
        vol_change = vol.diff()
        
        # 计算趋势因子
        trend_factor = price_change * vol_change
        
        return trend_factor.alias(self.factor_id)