NewStock/main/factor/polars_factors.py

"""
Polars因子主入口 - 整合所有Polars-based因子计算
提供统一的接口来应用所有类别的因子
"""

import polars as pl
from typing import Dict, List, Optional, Any
import logging

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


# 因子类别映射
FACTOR_CATEGORIES = {
    'money_flow': '资金流因子',
    'chip': '筹码分布因子', 
    'volatility': '波动率因子',
    'volume': '成交量因子',
    'technical': '技术指标因子',
    'sentiment': '情绪因子',
    'momentum': '动量因子',
    'complex': '复杂组合因子'
}


def apply_money_flow_factors(df: pl.DataFrame, operators: List = None) -> pl.DataFrame:
    """应用资金流因子"""
    try:
        from polars_money_flow_factors import apply_money_flow_factors as _apply_money_flow
        return _apply_money_flow(df, operators)
    except ImportError as e:
        logger.warning(f"无法导入资金流因子模块: {e}")
        return df


def apply_chip_distribution_factors(df: pl.DataFrame, operators: List = None) -> pl.DataFrame:
    """应用筹码分布因子"""
    try:
        from polars_chip_factors import apply_chip_distribution_factors as _apply_chip
        return _apply_chip(df, operators)
    except ImportError as e:
        logger.warning(f"无法导入筹码分布因子模块: {e}")
        return df


def apply_volatility_factors(df: pl.DataFrame, operators: List = None) -> pl.DataFrame:
    """应用波动率因子"""
    try:
        from polars_volatility_factors import apply_volatility_factors as _apply_volatility
        return _apply_volatility(df, operators)
    except ImportError as e:
        logger.warning(f"无法导入波动率因子模块: {e}")
        return df


def apply_volume_factors(df: pl.DataFrame, operators: List = None) -> pl.DataFrame:
    """应用成交量因子"""
    try:
        from polars_volume_factors import apply_volume_factors as _apply_volume
        return _apply_volume(df, operators)
    except ImportError as e:
        logger.warning(f"无法导入成交量因子模块: {e}")
        return df


def apply_technical_factors(df: pl.DataFrame, operators: List = None) -> pl.DataFrame:
    """应用技术指标因子"""
    try:
        from polars_technical_factors import apply_technical_factors as _apply_technical
        return _apply_technical(df, operators)
    except ImportError as e:
        logger.warning(f"无法导入技术指标因子模块: {e}")
        return df


def apply_sentiment_factors(df: pl.DataFrame, operators: List = None) -> pl.DataFrame:
    """应用情绪因子"""
    try:
        from polars_sentiment_factors import apply_sentiment_factors as _apply_sentiment
        return _apply_sentiment(df, operators)
    except ImportError as e:
        logger.warning(f"无法导入情绪因子模块: {e}")
        return df


def apply_momentum_factors(df: pl.DataFrame, operators: List = None) -> pl.DataFrame:
    """应用动量因子"""
    try:
        from polars_momentum_factors import apply_momentum_factors as _apply_momentum
        return _apply_momentum(df, operators)
    except ImportError as e:
        logger.warning(f"无法导入动量因子模块: {e}")
        return df


def apply_complex_factors(df: pl.DataFrame, operators: List = None) -> pl.DataFrame:
    """应用复杂组合因子"""
    try:
        from polars_complex_factors import apply_complex_factors as _apply_complex
        return _apply_complex(df, operators)
    except ImportError as e:
        logger.warning(f"无法导入复杂组合因子模块: {e}")
        return df


def apply_all_factors(df: pl.DataFrame, 
                     factor_categories: List[str] = None,
                     exclude_categories: List[str] = None) -> pl.DataFrame:
    """
    应用所有类别的因子
    
    Args:
        df: 输入的Polars DataFrame，必须包含必需的列
        factor_categories: 要应用的因子类别列表，如果为None则应用所有类别
        exclude_categories: 要排除的因子类别列表
    
    Returns:
        添加了所有因子的DataFrame
    """
    if factor_categories is None:
        factor_categories = list(FACTOR_CATEGORIES.keys())
    
    if exclude_categories:
        factor_categories = [cat for cat in factor_categories if cat not in exclude_categories]
    
    logger.info(f"开始应用因子类别: {factor_categories}")
    
    result_df = df
    total_factors = 0
    
    # 因子类别到函数的映射
    factor_functions = {
        'money_flow': apply_money_flow_factors,
        'chip': apply_chip_distribution_factors,
        'volatility': apply_volatility_factors,
        'volume': apply_volume_factors,
        'technical': apply_technical_factors,
        'sentiment': apply_sentiment_factors,
        'momentum': apply_momentum_factors,
        'complex': apply_complex_factors
    }
    
    for category in factor_categories:
        if category not in factor_functions:
            logger.warning(f"未知的因子类别: {category}")
            continue
        
        logger.info(f"应用{FACTOR_CATEGORIES[category]}...")
        
        try:
            before_cols = len(result_df.columns)
            result_df = factor_functions[category](result_df)
            after_cols = len(result_df.columns)
            new_factors = after_cols - before_cols
            
            logger.info(f"{FACTOR_CATEGORIES[category]}应用完成，新增{new_factors}个因子")
            total_factors += new_factors
            
        except Exception as e:
            logger.error(f"应用{FACTOR_CATEGORIES[category]}时出错: {e}")
            continue
    
    logger.info(f"因子应用完成，总共新增{total_factors}个因子")
    return result_df


def get_factor_info() -> Dict[str, Any]:
    """
    获取因子信息
    
    Returns:
        包含因子类别信息的字典
    """
    return {
        'categories': FACTOR_CATEGORIES,
        'total_categories': len(FACTOR_CATEGORIES),
        'category_descriptions': list(FACTOR_CATEGORIES.values())
    }


def validate_required_columns(df: pl.DataFrame, factor_categories: List[str] = None) -> Dict[str, List[str]]:
    """
    验证DataFrame是否包含必需的列
    
    Args:
        df: 输入的Polars DataFrame
        factor_categories: 要验证的因子类别列表
    
    Returns:
        包含缺失列信息的字典
    """
    if factor_categories is None:
        factor_categories = list(FACTOR_CATEGORIES.keys())
    
    missing_columns = {}
    
    # 基础必需列
    base_required = ['ts_code', 'trade_date']
    missing_base = [col for col in base_required if col not in df.columns]
    if missing_base:
        missing_columns['base'] = missing_base
    
    # 各因子类别的必需列
    category_requirements = {
        'money_flow': ['buy_lg_vol', 'buy_elg_vol', 'sell_lg_vol', 'sell_elg_vol', 'vol'],
        'chip': ['cost_95pct', 'cost_85pct', 'cost_50pct', 'cost_15pct', 'cost_5pct', 
                'winner_rate', 'weight_avg', 'close'],
        'volatility': ['pct_chg'],
        'volume': ['vol', 'turnover_rate', 'volume_ratio', 'amount'],
        'technical': ['open', 'high', 'low', 'close', 'vol'],
        'sentiment': ['pct_chg', 'vol', 'volume_ratio'],
        'momentum': ['close', 'turnover_rate'],
        'complex': ['close', 'vol', 'pct_chg', 'turnover_rate', 'winner_rate']
    }
    
    for category in factor_categories:
        if category in category_requirements:
            required_cols = category_requirements[category]
            missing_cols = [col for col in required_cols if col not in df.columns]
            if missing_cols:
                missing_columns[category] = missing_cols
    
    return missing_columns


# 向后兼容的函数名
apply_factors = apply_all_factors


if __name__ == "__main__":
    # 测试代码
    print("Polars因子系统已加载")
    print("可用的因子类别:")
    for key, description in FACTOR_CATEGORIES.items():
        print(f"  {key}: {description}")