import polars as pl

from ldnode.calculator.schemas import HoldWeight, BacktestParameters


def planning_weight_column(
    frame: pl.LazyFrame | pl.DataFrame,
    parameters: BacktestParameters,
    with_exclude: bool = False,
) -> pl.LazyFrame:
    """ 规划持仓权重 """
    lf = frame.lazy() if isinstance(frame, pl.DataFrame) else frame
    # 选出排名内的标的
    _between_rank_expr = pl.col('rank').is_between(
        lower_bound=parameters['hold_method']['hold_range_value']['start'],
        upper_bound=parameters['hold_method']['hold_range_value']['end']
    )
    # 根据持仓方式计算权重
    if parameters['hold_method']['hold_weight'] == HoldWeight.AVG:
        # 等权持有
        weight_expr = (
            pl.when(_between_rank_expr).then(
                pl.lit(1).truediv(
                    pl.col('code').filter(_between_rank_expr).len().over('trade_date')
                )
            ).otherwise(None)
        ) if with_exclude else (
            pl.lit(1).truediv(
                pl.len().over('trade_date')
            )
        )
    elif parameters['hold_method']['hold_weight'] == HoldWeight.CLOSE:
        # 按价格比例持有
        weight_expr = (
            pl.when(_between_rank_expr).then(
                pl.col('close').truediv(
                    pl.col('close').filter(_between_rank_expr).sum().over('trade_date')
                )
            ).otherwise(None)
        ) if with_exclude else (
            pl.col('close').truediv(
                pl.col('close').sum().over('trade_date')
            )
        )
    elif parameters['hold_method']['hold_weight'] == HoldWeight.AMOUNT:
        # 按成交额比例持有
        weight_expr = (
            pl.when(_between_rank_expr).then(
                pl.col('amount').truediv(
                    pl.col('amount').filter(_between_rank_expr).sum().over('trade_date')
                ) 
            ).otherwise(None)
        ) if with_exclude else (
            pl.col('amount').truediv(
                pl.col('amount').sum().over('trade_date')
            )
        )

    elif parameters['hold_method']['hold_weight'] == HoldWeight.NOTHOLD:
        weight_expr = pl.lit(.0)
    else:
        raise ValueError(f"Invalid hold weight: {parameters['hold_method']['hold_weight']}")
    
    # 添加权重列（向量回测专用，列名为 weight，与 vector_backtest.py 保持一致）
    return (
        lf.with_columns(
            weight_expr.alias('weight'),
        )
    )
