"""
技术指标因子 - 使用Polars实现
包含ATR、OBV、RSI、EMA等技术指标相关因子计算
"""

import polars as pl
import numpy as np
from typing import Dict, List, Optional, Any
from operator_framework import StockWiseOperator, OperatorConfig
import talib


class ATROperator(StockWiseOperator):
    """ATR算子"""
    
    def __init__(self, period: int = 14):
        config = OperatorConfig(
            name=f"atr_{period}",
            description=f"{period}日ATR",
            required_columns=['high', 'low', 'close'],
            output_columns=[f'atr_{period}'],
            parameters={'period': period}
        )
        super().__init__(config)
        self.period = period
    
    def apply_stock(self, stock_df: pl.DataFrame, **kwargs) -> pl.DataFrame:
        """计算ATR"""
        # 使用TA-Lib计算ATR
        atr_values = talib.ATR(
            stock_df['high'].to_numpy(),
            stock_df['low'].to_numpy(), 
            stock_df['close'].to_numpy(),
            timeperiod=self.period
        )
        
        return stock_df.with_columns(pl.Series(atr_values).alias(f'atr_{self.period}'))


class OBVOperator(StockWiseOperator):
    """OBV算子"""
    
    def __init__(self):
        config = OperatorConfig(
            name="obv",
            description="OBV能量潮",
            required_columns=['close', 'vol'],
            output_columns=['obv'],
            parameters={}
        )
        super().__init__(config)
    
    def apply_stock(self, stock_df: pl.DataFrame, **kwargs) -> pl.DataFrame:
        """计算OBV"""
        # 使用TA-Lib计算OBV
        obv_values = talib.OBV(
            stock_df['close'].to_numpy(),
            stock_df['vol'].to_numpy()
        )
        
        return stock_df.with_columns(pl.Series(obv_values).alias('obv'))


class OBVMAOperator(StockWiseOperator):
    """OBV均线算子"""
    
    def __init__(self, period: int = 6):
        config = OperatorConfig(
            name=f"obv_ma_{period}",
            description=f"{period}日OBV均线",
            required_columns=['obv'],
            output_columns=[f'maobv_{period}'],
            parameters={'period': period}
        )
        super().__init__(config)
        self.period = period
    
    def apply_stock(self, stock_df: pl.DataFrame, **kwargs) -> pl.DataFrame:
        """计算OBV均线"""
        # 使用TA-Lib计算SMA
        ma_values = talib.SMA(
            stock_df['obv'].to_numpy(),
            timeperiod=self.period
        )
        
        return stock_df.with_columns(pl.Series(ma_values).alias(f'maobv_{self.period}'))


class RSIOperator(StockWiseOperator):
    """RSI算子"""
    
    def __init__(self, period: int = 3):
        config = OperatorConfig(
            name=f"rsi_{period}",
            description=f"{period}日RSI",
            required_columns=['close'],
            output_columns=[f'rsi_{period}'],
            parameters={'period': period}
        )
        super().__init__(config)
        self.period = period
    
    def apply_stock(self, stock_df: pl.DataFrame, **kwargs) -> pl.DataFrame:
        """计算RSI"""
        # 使用TA-Lib计算RSI
        rsi_values = talib.RSI(
            stock_df['close'].to_numpy(),
            timeperiod=self.period
        )
        
        return stock_df.with_columns(pl.Series(rsi_values).alias(f'rsi_{self.period}'))


class EMAOperator(StockWiseOperator):
    """EMA算子"""
    
    def __init__(self, period: int):
        config = OperatorConfig(
            name=f"ema_{period}",
            description=f"{period}日EMA",
            required_columns=['close'],
            output_columns=[f'_ema_{period}'],
            parameters={'period': period}
        )
        super().__init__(config)
        self.period = period
    
    def apply_stock(self, stock_df: pl.DataFrame, **kwargs) -> pl.DataFrame:
        """计算EMA"""
        # 使用TA-Lib计算EMA
        ema_values = talib.EMA(
            stock_df['close'].to_numpy(),
            timeperiod=self.period
        )
        
        return stock_df.with_columns(pl.Series(ema_values).alias(f'_ema_{self.period}'))


class ReturnOperator(StockWiseOperator):
    """收益率算子"""
    
    def __init__(self, period: int):
        config = OperatorConfig(
            name=f"return_{period}",
            description=f"{period}日收益率",
            required_columns=['close'],
            output_columns=[f'return_{period}'],
            parameters={'period': period}
        )
        super().__init__(config)
        self.period = period
    
    def apply_stock(self, stock_df: pl.DataFrame, **kwargs) -> pl.DataFrame:
        """计算收益率"""
        # 计算收益率
        ret = pl.col('close').pct_change(self.period)
        
        return stock_df.with_columns(ret.alias(f'return_{self.period}'))


class ActivityFactorOperator(StockWiseOperator):
    """活跃度因子算子"""
    
    def __init__(self, period: int, scale: float):
        config = OperatorConfig(
            name=f"act_factor_{period}",
            description=f"{period}日活跃度因子",
            required_columns=[f'_ema_{period}'],
            output_columns=[f'act_factor{period}'],
            parameters={'period': period, 'scale': scale}
        )
        super().__init__(config)
        self.period = period
        self.scale = scale
    
    def apply_stock(self, stock_df: pl.DataFrame, **kwargs) -> pl.DataFrame:
        """计算活跃度因子"""
        # 计算EMA变化率
        ema_change = (pl.col(f'_ema_{self.period}') / pl.col(f'_ema_{self.period}').shift(1) - 1) * 100
        
        # 计算活跃度因子
        activity_factor = (ema_change * 57.3 / self.scale).arctan()
        
        return stock_df.with_columns(activity_factor.alias(f'act_factor{self.period}'))


class ActivityFactor5Operator(StockWiseOperator):
    """活跃度因子5算子"""
    
    def __init__(self):
        config = OperatorConfig(
            name="act_factor_5",
            description="5日活跃度因子",
            required_columns=['_ema_5'],
            output_columns=['act_factor1'],
            parameters={}
        )
        super().__init__(config)
    
    def apply_stock(self, stock_df: pl.DataFrame, **kwargs) -> pl.DataFrame:
        """计算5日活跃度因子"""
        # 计算EMA变化率
        ema_change = (pl.col('_ema_5') / pl.col('_ema_5').shift(1) - 1) * 100
        
        # 计算活跃度因子
        activity_factor = (ema_change * 57.3 / 50).arctan()
        
        return stock_df.with_columns(activity_factor.alias('act_factor1'))


class ActivityFactor13Operator(StockWiseOperator):
    """活跃度因子13算子"""
    
    def __init__(self):
        config = OperatorConfig(
            name="act_factor_13",
            description="13日活跃度因子",
            required_columns=['_ema_13'],
            output_columns=['act_factor2'],
            parameters={}
        )
        super().__init__(config)
    
    def apply_stock(self, stock_df: pl.DataFrame, **kwargs) -> pl.DataFrame:
        """计算13日活跃度因子"""
        # 计算EMA变化率
        ema_change = (pl.col('_ema_13') / pl.col('_ema_13').shift(1) - 1) * 100
        
        # 计算活跃度因子
        activity_factor = (ema_change * 57.3 / 40).arctan()
        
        return stock_df.with_columns(activity_factor.alias('act_factor2'))


class ActivityFactor20Operator(StockWiseOperator):
    """活跃度因子20算子"""
    
    def __init__(self):
        config = OperatorConfig(
            name="act_factor_20",
            description="20日活跃度因子",
            required_columns=['_ema_20'],
            output_columns=['act_factor3'],
            parameters={}
        )
        super().__init__(config)
    
    def apply_stock(self, stock_df: pl.DataFrame, **kwargs) -> pl.DataFrame:
        """计算20日活跃度因子"""
        # 计算EMA变化率
        ema_change = (pl.col('_ema_20') / pl.col('_ema_20').shift(1) - 1) * 100
        
        # 计算活跃度因子
        activity_factor = (ema_change * 57.3 / 21).arctan()
        
        return stock_df.with_columns(activity_factor.alias('act_factor3'))


class ActivityFactor60Operator(StockWiseOperator):
    """活跃度因子60算子"""
    
    def __init__(self):
        config = OperatorConfig(
            name="act_factor_60",
            description="60日活跃度因子",
            required_columns=['_ema_60'],
            output_columns=['act_factor4'],
            parameters={}
        )
        super().__init__(config)
    
    def apply_stock(self, stock_df: pl.DataFrame, **kwargs) -> pl.DataFrame:
        """计算60日活跃度因子"""
        # 计算EMA变化率
        ema_change = (pl.col('_ema_60') / pl.col('_ema_60').shift(1) - 1) * 100
        
        # 计算活跃度因子
        activity_factor = (ema_change * 57.3 / 10).arctan()
        
        return stock_df.with_columns(activity_factor.alias('act_factor4'))


class ActivityFactor5and6Operator(StockWiseOperator):
    """活跃度因子5和6算子"""
    
    def __init__(self):
        config = OperatorConfig(
            name="act_factor_5_6",
            description="活跃度因子5和6",
            required_columns=['act_factor1', 'act_factor2'],
            output_columns=['act_factor5', 'act_factor6'],
            parameters={}
        )
        super().__init__(config)
    
    def apply_stock(self, stock_df: pl.DataFrame, **kwargs) -> pl.DataFrame:
        """计算活跃度因子5和6"""
        # 计算因子5
        factor5 = pl.col('act_factor1') + pl.col('act_factor2') + pl.col('act_factor3') + pl.col('act_factor4')
        
        # 计算因子6
        numerator = pl.col('act_factor1') - pl.col('act_factor2')
        denominator = (pl.col('act_factor1').pow(2) + pl.col('act_factor2').pow(2)).sqrt()
        factor6 = numerator / (denominator + 1e-8)  # 避免除零
        
        return stock_df.with_columns([
            factor5.alias('act_factor5'),
            factor6.alias('act_factor6')
        ])


class Alpha003Operator(StockWiseOperator):
    """Alpha003算子"""
    
    def __init__(self):
        config = OperatorConfig(
            name="alpha_003",
            description="Alpha003因子",
            required_columns=['open', 'close', 'high', 'low'],
            output_columns=['alpha_003'],
            parameters={}
        )
        super().__init__(config)
    
    def apply_stock(self, stock_df: pl.DataFrame, **kwargs) -> pl.DataFrame:
        """计算Alpha003"""
        # 计算因子
        alpha_003 = pl.when(pl.col('high') != pl.col('low')) \
            .then((pl.col('close') - pl.col('open')) / (pl.col('high') - pl.col('low'))) \
            .otherwise(0)
        
        return stock_df.with_columns(alpha_003.alias('alpha_003'))


class Alpha007Operator(StockWiseOperator):
    """Alpha007算子"""
    
    def __init__(self):
        config = OperatorConfig(
            name="alpha_007",
            description="Alpha007因子",
            required_columns=['close', 'vol'],
            output_columns=['alpha_007'],
            parameters={}
        )
        super().__init__(config)
    
    def apply_stock(self, stock_df: pl.DataFrame, **kwargs) -> pl.DataFrame:
        """计算Alpha007"""
        # 计算5日相关性
        corr_5 = pl.col('close').rolling_corr(pl.col('vol'), window=5)
        
        return stock_df.with_columns(corr_5.alias('alpha_007'))


class Alpha013Operator(StockWiseOperator):
    """Alpha013算子"""
    
    def __init__(self):
        config = OperatorConfig(
            name="alpha_013",
            description="Alpha013因子",
            required_columns=['close'],
            output_columns=['alpha_013'],
            parameters={}
        )
        super().__init__(config)
    
    def apply_stock(self, stock_df: pl.DataFrame, **kwargs) -> pl.DataFrame:
        """计算Alpha013"""
        # 计算5日和20日和
        sum_5 = pl.col('close').rolling_sum(window=5)
        sum_20 = pl.col('close').rolling_sum(window=20)
        
        # 计算因子
        alpha_013 = sum_5 - sum_20
        
        return stock_df.with_columns(alpha_013.alias('alpha_013'))


class Alpha022Operator(StockWiseOperator):
    """Alpha022算子"""
    
    def __init__(self):
        config = OperatorConfig(
            name="alpha_022",
            description="Alpha022改进因子",
            required_columns=['high', 'low', 'close', 'vol'],
            output_columns=['alpha_22_improved'],
            parameters={}
        )
        super().__init__(config)
    
    def apply_stock(self, stock_df: pl.DataFrame, **kwargs) -> pl.DataFrame:
        """计算Alpha022改进因子"""
        # 计算滚动协方差
        cov_5 = pl.col('high').rolling_cov(pl.col('vol'), window=5)
        
        # 计算协方差差分
        delta_cov = cov_5.diff(5)
        
        # 计算收盘价标准差
        std_close = pl.col('close').rolling_std(window=20)
        
        # 计算标准差排名 (简化版)
        rank_std = std_close
        
        # 计算最终因子
        alpha_22 = -1 * delta_cov * rank_std
        
        return stock_df.with_columns(alpha_22.alias('alpha_22_improved'))


class BBIRatioOperator(StockWiseOperator):
    """BBI比率算子"""
    
    def __init__(self):
        config = OperatorConfig(
            name="bbi_ratio",
            description="BBI比率因子",
            required_columns=['close'],
            output_columns=['bbi_ratio_factor'],
            parameters={}
        )
        super().__init__(config)
    
    def apply_stock(self, stock_df: pl.DataFrame, **kwargs) -> pl.DataFrame:
        """计算BBI比率"""
        # 计算不同周期的SMA
        sma3 = pl.col('close').rolling_mean(window=3)
        sma6 = pl.col('close').rolling_mean(window=6)
        sma12 = pl.col('close').rolling_mean(window=12)
        sma24 = pl.col('close').rolling_mean(window=24)
        
        # 计算BBI
        bbi = (sma3 + sma6 + sma12 + sma24) / 4
        
        # 计算比率
        bbi_ratio = bbi / pl.col('close')
        
        return stock_df.with_columns(bbi_ratio.alias('bbi_ratio_factor'))


# 技术指标因子集合
TECHNICAL_OPERATORS = [
    ATROperator(14),
    ATROperator(6),
    OBVOperator(),
    OBVMAOperator(6),
    RSIOperator(3),
    EMAOperator(5),
    EMAOperator(13),
    EMAOperator(20),
    EMAOperator(60),
    ReturnOperator(5),
    ReturnOperator(20),
    ActivityFactor5Operator(),
    ActivityFactor13Operator(),
    ActivityFactor20Operator(),
    ActivityFactor60Operator(),
    ActivityFactor5and6Operator(),
    Alpha003Operator(),
    Alpha007Operator(),
    Alpha013Operator(),
    Alpha022Operator(),
    BBIRatioOperator(),
]


def apply_technical_factors(df: pl.DataFrame, operators: List = None) -> pl.DataFrame:
    """
    应用所有技术指标因子
    
    Args:
        df: 输入的Polars DataFrame
        operators: 要应用的算子列表，如果为None则使用默认列表
    
    Returns:
        添加了技术指标因子的DataFrame
    """
    if operators is None:
        operators = TECHNICAL_OPERATORS
    
    result_df = df
    for operator in operators:
        result_df = operator(result_df)
    
    return result_df