from typing import Callable, Optional, Type

import polars as pl
from ldnode.calculator.schemas import (
    BuiltinScope,
    HoldRangeType,
    NeutralizationMethod,
    ProcessMissingDataMethod,
    ScreenParameters,
)

from .computed_factors import planning_computed_factors
from .exclude_factors import planning_exclude_factors
from .missing_data import planning_missing_data_process
from .neutralization import planning_neutralization_process
from .pct_factors import planning_pct_factors
from .rank_factors import planning_rank_factors
from .score_factors import planning_score_factors
from .security_pools import planning_security_pools


def planning_screen_lf(
    lf: pl.LazyFrame,
    parameters: ScreenParameters,
    builtin_factors: list[str],
    filter_exprs: list[pl.Expr],
    BuiltinScopes: Type[BuiltinScope],
    ProcessMissingDataMethods: Type[ProcessMissingDataMethod],
    NeutralizationMethods: Type[NeutralizationMethod],
    filter_fn: Optional[Callable[[pl.LazyFrame], pl.LazyFrame]] = None,
    keep_exclude: bool = True,
    not_rank: bool = False,
) -> pl.LazyFrame:
    """规划选债/选股截面"""
    # 规划自定义因子计算
    lf = lf.pipe(
        planning_computed_factors,
        builtin_factors=builtin_factors,
        computed_factors=parameters.get("computed_factors"),
        exclude_suspension=False,
    )
    # 按照规则排除截面数据
    if filter_fn:
        lf = filter_fn(lf)
    lf = (
        lf
        # 规划百分比因子
        .pipe(
            planning_pct_factors,
            screen_parameters=parameters,
        )
        # 处理标的池
        .pipe(
            planning_security_pools,
            BuiltinScopes=BuiltinScopes,
            security_pools=parameters.get("pools"),
        )
        # 规划排除因子
        .pipe(
            planning_exclude_factors,
            filter_factors=parameters.get("filter_factors"),
            filter_exprs=filter_exprs,
            keep_exclude=keep_exclude,
        )
        # 规划缺失值处理
        .pipe(
            planning_missing_data_process,
            rank_factors=parameters.get("rank_factors"),
            ProcessMethods=ProcessMissingDataMethods,
        )
        # 规划中性化处理
        .pipe(
            planning_neutralization_process,
            rank_factors=parameters.get("rank_factors"),
            NeutralizationMethods=NeutralizationMethods,
        )
        # 规划得分计算
        .pipe(
            planning_score_factors,
            rank_factors=parameters.get("rank_factors"),
            with_exclude=keep_exclude,
        )
    )
    if not not_rank:
        lf = lf.pipe(
            planning_rank_factors,
            pct=parameters["hold_method"]["hold_range_type"] == HoldRangeType.PCT,
            more_by=parameters.get("more_by"),
        )
    return lf
