import polars as pl
import polars.selectors as cs

from ldnode.calculator.schemas import RankFactor


def _gen_score_expr(
    rank_factor: RankFactor[str, str],
    with_exclude: bool = False,
) -> list[pl.Expr]:
    """ 生成因子打分表达式 """
    factor = f"{rank_factor['factor']}_neutralization" if rank_factor.get('neutralization_factor') else rank_factor['factor']
    if with_exclude:
        return (
            pl.when(pl.col('_exclude') | ~pl.col('_isin_pool'))
            .then(None)  # 被排除或者不在标的池内的标的得分设为 None
            .otherwise(pl.col(factor))
            .rank(method='average', descending=not rank_factor['ascending'])
            .over('trade_date')  # 按交易日期分组排名
            .mul(rank_factor['weight'])  # 乘以权重
            .alias(f'_{factor}_score')  # 使用原始因子名命名得分列
        )
    else:
        return (
            pl.col(factor)
            .rank(method='average', descending=not rank_factor['ascending'])
            .over('trade_date')
            .mul(rank_factor['weight'])
            .alias(f'_{factor}_score')
        )
        
    
def _gen_score_exprs(
    rank_factors: list[RankFactor[str, str]],
    with_exclude: bool = False,
) -> list[pl.Expr]:
    """ 生成因子打分表达式 """
    return [
        _gen_score_expr(rank_factor, with_exclude) for rank_factor in rank_factors
    ]


def planning_score_factors(
    frame: pl.LazyFrame | pl.DataFrame,
    rank_factors: list[RankFactor[str, str]],
    with_exclude: bool = False,
) -> pl.LazyFrame:
    """ 规划因子打分计算 """
    lf = frame.lazy() if isinstance(frame, pl.DataFrame) else frame
    if with_exclude:
        return (
            lf
            .with_columns(_gen_score_exprs(rank_factors, with_exclude))  # 添加各因子的得分列
            .with_columns(
                # 计算总分：如果标的被排除或者不在标的池内，总分为 None；否则对所有得分求和
                pl.when(pl.col('_exclude') | ~pl.col('_isin_pool'))
                .then(None)
                .otherwise(
                    pl.sum_horizontal(
                        cs.starts_with('_') & cs.ends_with('_score'), 
                        ignore_nulls=True  # 忽略空值
                    )
                )
                .alias('score')
            )
        )
    else:
        return (
            lf
            .with_columns(_gen_score_exprs(rank_factors, with_exclude))
            .with_columns(
                # 直接对所有得分列求和得到总分
                pl.sum_horizontal(
                    cs.starts_with('_') & cs.ends_with('_score'), 
                    ignore_nulls=True
                )
                .alias('score')
            )
        )
    