import polars as pl
from polars import col

from ldnode.calculator.schemas import (
    RebalanceType,
    BacktestParameters,
)


def calc_rebalance_days(frame: pl.LazyFrame | pl.DataFrame, parameters: BacktestParameters) -> pl.DataFrame:
    """
    计算再平衡日
    :param df: 交易日历df，包含trade_date和index列
    :rtype: pl.LazyFrame | pl.DataFrame
    :param parameters: 策略配置
    :rtype: BaseStrategyConfig
    :return: 交易日历df，包含trade_date和if_rebalance列
    :rtype: pl.DataFrame
    """
    # 确保df为 LazyFrame
    lf = frame.lazy() if isinstance(frame, pl.DataFrame) else frame
    rebalance_type = parameters['rebalance_method']['rebalance_type']
    rebalance_value = parameters['rebalance_method']['rebalance_value']
    # 每若干个交易日调仓
    if rebalance_type == RebalanceType.TRADEDAY:
        return lf.with_columns(
            (col('index').mod(rebalance_value) == 0).alias('if_rebalance')
        ).select('trade_date', 'if_rebalance').collect()

    # 每星期调仓
    elif rebalance_type == RebalanceType.WEEKDAY:
        return lf.with_columns(
            (col('trade_date').dt.weekday() == rebalance_value).alias('if_rebalance')
        ).select('trade_date', 'if_rebalance').collect()
    else:
        raise ValueError(f"Invalid rebalance type: {rebalance_type}")
