import polars as pl

from ldnode.calculator.schemas import HoldMethod, HoldWeight


def planning_screen_weight(
    lf: pl.LazyFrame,
    hold_method: HoldMethod,
) -> pl.LazyFrame:
    """ 规划选标权重 """
    match hold_method['hold_weight']:
        case HoldWeight.AVG:
            # 等权持有
            weight_expr = pl.lit(1)
        case HoldWeight.CLOSE:
            # 按价格比例持有
            weight_expr = pl.col('close')
        case HoldWeight.AMOUNT:
            # 按成交额比例持有
            weight_expr = pl.col('amount')
        case HoldWeight.NOTHOLD:
            # 不持有
            weight_expr = pl.lit(.0)
        case _:
            raise ValueError(f"Invalid hold weight: {hold_method['hold_weight']}")
    return lf.with_columns(
        weight_expr.alias('_weight')
    )
