from typing import Type

import polars as pl

from ldnode.calculator.schemas import (
    RankFactor,
    ProcessMissingDataMethod,
)
        

def _gen_missing_data_exprs(
    rank_factors: list[RankFactor[str, str]],
    ProcessMethods: Type[ProcessMissingDataMethod],
) -> list[tuple[str, pl.Expr]]:
    return  [
        (
            rank_factor['factor'],
            ProcessMethods(rank_factor['process_missing_data_method']) \
                .process(rank_factor['factor']).alias(rank_factor['factor'])
        ) for rank_factor in rank_factors
        if rank_factor.get('process_missing_data_method') is not None \
            and rank_factor['process_missing_data_method'] in ProcessMethods
    ]

def planning_missing_data_process(
    frame: pl.LazyFrame | pl.DataFrame,
    rank_factors: list[RankFactor[str, str]],
    ProcessMethods: Type[ProcessMissingDataMethod],
) -> pl.LazyFrame:
    """ 规划缺失值处理 """
    lf = frame.lazy() if isinstance(frame, pl.DataFrame) else frame
    # 检查是否有因子需要缺失值处理
    if all(rank_factor.get('process_missing_data_method') is None for rank_factor in rank_factors):
        return lf
    # 保存停牌记录，后续恢复
    suspension_lf = lf.filter(pl.col('suspension'))
    lf = lf.filter(~pl.col('suspension'))
    # 获取缺失值处理表达式
    if name_and_exprs := _gen_missing_data_exprs(rank_factors, ProcessMethods):
        lf = lf.with_columns([expr for _, expr in name_and_exprs])
        # 为停牌数据中的所有已处理因子设置为None
        if processed_factors := [name for name, _ in name_and_exprs]:
            suspension_columns = [pl.lit(None).alias(factor_name) for factor_name in processed_factors]
            suspension_lf = suspension_lf.with_columns(suspension_columns)
    # 合并停牌和非停牌数据
    lf = pl.concat([lf, suspension_lf]).sort('trade_date')
    return lf
