from typing import Optional

import polars as pl

from ldnode.calculator.schemas import SortOrder


def planning_rank_factors(
    frame: pl.LazyFrame | pl.DataFrame,
    pct: bool = False,
    more_by: Optional[list[tuple[str, SortOrder]]] = None,
) -> pl.LazyFrame:
    """ 规划因子排名计算 """
    lf = frame.lazy() if isinstance(frame, pl.DataFrame) else frame
    # 处理more_by默认值
    more_by = more_by or [('close', SortOrder.ASCENDING)]
    # 构建排序表达式列表 - 确保score在前面
    sort_columns = ['trade_date', 'score']  # 首先按交易日期，然后按score
    sort_flags = [False, True]  # trade_date 升序, score 降序
    # 添加more_by字段的表达式
    for field, sort_order in more_by:
        sort_columns.append(field)
        sort_flags.append(sort_order == SortOrder.DESCENDING)
    # 计算排名表达式
    if pct:
        rank_expr = (
            pl.col('score')
            .rank(method='average', descending=True)
            .over('trade_date')
            .truediv(
                pl.col('score').count().over('trade_date')
            )
        )
    else:
        rank_expr = (
            pl.col('score')
            .rank(method='ordinal', descending=True)
            .over('trade_date')
        )
    return (
        lf
        .sort(sort_columns, descending=sort_flags)  # 使用表达式列表，确保排序包含score
        .with_columns(rank_expr.alias('rank'))
    )
    