import polars as pl
from ldnode.calculator.schemas import (
    ExcludeConditionGroup,
    ExcludeFactor,
    ScreenParameters,
    SecurityPoolCustomizeType,
    SecurityPoolType,
)


def planning_pct_factors(
    frame: pl.LazyFrame | pl.DataFrame,
    screen_parameters: ScreenParameters,
) -> pl.LazyFrame:
    """规划动态股票池所需的因子百分比排行计算"""
    lf = frame.lazy() if isinstance(frame, pl.DataFrame) else frame
    filter_factors: list[ExcludeFactor] = []
    # 仅处理动态股票池 pool_filter_factors 所需的 _pct
    for pool in screen_parameters.get("pools", []) or []:
        if (
            pool["pool_type"] == SecurityPoolType.CUSTOMIZED
            and pool.get("pool_customize_type")
            == SecurityPoolCustomizeType.DYNAMIC_BY_FACTOR
            and pool.get("pool_filter_factors")
        ):
            filter_factors.extend(pool.get("pool_filter_factors"))
    factor_pct_exprs = [
        (
            pl.col(factor)
            .rank(method="average", descending=False)
            .over("trade_date")
            .truediv(
                pl.col(factor)
                .count()
                .over("trade_date")
                .sub(pl.col(factor).null_count().over("trade_date"))
            )
            .alias(f"{factor}_pct")
        )
        for factor in set(
            [
                filter_factor["factor"]
                for filter_factor in filter_factors
                if "_pct" in filter_factor["operator"]
            ]
        )
    ]
    if factor_pct_exprs:
        lf = lf.with_columns(factor_pct_exprs)
    return lf
