import polars as pl

from ldnode.calculator.schemas import HoldRangeType


def planning_mod_rank(
    frame: pl.LazyFrame | pl.DataFrame,
    threshold: int,
    hold_range_type: HoldRangeType,
    position_lf: pl.LazyFrame
) -> pl.LazyFrame:
    """ 规划修正排名 """
    lf = frame.lazy() if isinstance(frame, pl.DataFrame) else frame
    lf = (
        # 合并持仓中的volume
        lf.join(position_lf.select('code', 'volume'), on='code', how='left')
            .with_columns(
                pl.col('volume').fill_null(.0).alias('volume')
            )
            # 记录原排名
            .rename({'rank': 'ori_rank'})
    )
    # 计算修正排名表达式
    mod_rank_expr = (
        pl.when(pl.col('volume') > 0)
        .then(pl.col('ori_rank') - threshold)
        .otherwise(pl.col('ori_rank'))
        .alias('mod_rank')
    )
    # 按数量排名时
    match hold_range_type:
        case HoldRangeType.NUM:
            lf = (
                lf.with_columns(mod_rank_expr)
                    .sort(['mod_rank','volume'], descending=[False,True])
                    .with_columns(pl.col('mod_rank').rank(method='ordinal').alias('rank'))
            )
        case HoldRangeType.PCT:
            lf = lf.with_columns(mod_rank_expr.alias('rank'))
    return lf
