from datetime import date, datetime
from typing import Type

import polars as pl
from polars_event_backtest import event_backtest as rs_event_backtest

from ldnode import logger
from ldnode.calculator.plans.polars import (
    planning_mod_rank,
    planning_screen_lf,
)
from ldnode.calculator.schemas import (
    CB_BUILTIN_FACTORS,
    EPSILON,
    ETF_BUILTIN_FACTORS,
    STK_BUILTIN_FACTORS,
    BacktestParameters,
    BuiltinScope,
    CBBacktestParameters,
    CBBuiltinScope,
    CBNeutralizationMethod,
    CBProcessMissingDataMethod,
    ETFBacktestParameters,
    ETFBuiltinScope,
    ETFNeutralizationMethod,
    ETFProcessMissingDataMethod,
    HistoryPosition,
    HoldWeight,
    NeutralizationMethod,
    Portfolio,
    ProcessMissingDataMethod,
    STKBacktestParameters,
    STKBuiltinScope,
    STKNeutralizationMethod,
    STKProcessMissingDataMethod,
    StopType,
    gen_cb_filter_exprs,
    gen_etf_filter_exprs,
    gen_stk_filter_exprs,
)

from .calc_rebalance_days import calc_rebalance_days


def _rs_event_backtest(
    lf: pl.LazyFrame,
    parameters: BacktestParameters,
    builtin_factors: list[str],
    BuiltinScopes: Type[BuiltinScope],
    ProcessMissingDataMethods: Type[ProcessMissingDataMethod],
    NeutralizationMethods: Type[NeutralizationMethod],
    board_lot: int,
    filter_exprs: list[pl.Expr] = [],
    record_history: bool = True,
    # Issue #1：新增参数，控制是否在返回值中附带 planning_screen_lf 全列截面数据
    # factor_analysis handler 需要全因子列来计算 Rank IC
    return_screen_df: bool = False,
) -> tuple[pl.DataFrame, pl.DataFrame] | tuple[pl.DataFrame, pl.DataFrame, pl.DataFrame]:
    """事件驱动回测(Rust实现)"""
    # 交易日范围series
    trade_dates_df = (
        lf.filter(
            pl.col("trade_date").is_between(
                lower_bound=datetime.strptime(
                    parameters["date_range"]["start_date"], "%Y-%m-%d"
                ),
                upper_bound=datetime.strptime(
                    parameters["date_range"]["end_date"], "%Y-%m-%d"
                ),
            )
        )
        .select("trade_date")
        .unique()
        .sort("trade_date")
        .set_sorted("trade_date")  # 将列标记为已排序,加快未来的操作
    ).collect()
    filter_exprs.extend([pl.col("suspension"), pl.col("last_trade_day")])
    lf = (
        lf.pipe(
            planning_screen_lf,
            parameters=parameters,
            builtin_factors=builtin_factors,
            filter_exprs=filter_exprs,
            BuiltinScopes=BuiltinScopes,
            ProcessMissingDataMethods=ProcessMissingDataMethods,
            NeutralizationMethods=NeutralizationMethods,
            filter_fn=lambda x: x.filter(
                pl.col("trade_date").is_between(
                    lower_bound=datetime.strptime(
                        parameters["date_range"]["start_date"], "%Y-%m-%d"
                    ),
                    upper_bound=datetime.strptime(
                        parameters["date_range"]["end_date"], "%Y-%m-%d"
                    ),
                )
            ),
        ).sort(["trade_date", "rank"], descending=False)  # 同交易日内按rank排序
    )
    rebalance_days = calc_rebalance_days(
        frame=trade_dates_df.with_row_index(), parameters=parameters
    )
    # 事件循环所需列
    event_loop_need_columns = [
        "code",
        "name",
        "trade_date",
        "close",
        "stk_div",
        "cash_div",
        "rank",
        "limit",
        "suspension",
    ]
    # 添加持仓权重所需列
    if parameters["hold_method"]["hold_weight"] == HoldWeight.AMOUNT:
        event_loop_need_columns.append("amount")  # 添加成交量列
    # 添加止盈/止损所需列
    if parameters["stop_method"]["stop_type"] != StopType.NOTSTOP:
        _stop_trade_price_expr = (
            pl.col("pre_close") * (1 + parameters["stop_method"]["stop_value"])
        ).round(3)
        if parameters["stop_method"]["stop_type"] == StopType.TAKEPROFIT:
            # 若开盘价高于止盈价格，则止盈价格为开盘价
            lf = lf.with_columns(
                pl.when(pl.col("high") > _stop_trade_price_expr)
                .then(True)
                .otherwise(False)
                .alias("_stop_trade_flag"),
                pl.max_horizontal(_stop_trade_price_expr, pl.col("open")).alias(
                    "_stop_trade_price"
                ),
            )
        elif parameters["stop_method"]["stop_type"] == StopType.STOPLOSS:
            # 若开盘价低于止损价格，则止损价格为开盘价
            lf = lf.with_columns(
                pl.when(pl.col("low") < _stop_trade_price_expr)
                .then(True)
                .otherwise(False)
                .alias("_stop_trade_flag"),
                pl.min_horizontal(_stop_trade_price_expr, pl.col("open")).alias(
                    "_stop_trade_price"
                ),
            )
        event_loop_need_columns.extend(
            ["_stop_trade_price", "_stop_trade_flag"]
        )  # 添加止盈/止损价格列
    # factor_analysis 需要全因子列计算 Rank IC，其它场景只需事件循环列
    if return_screen_df:
        screen_df = lf.collect()
        event_df = screen_df.select(event_loop_need_columns)
    else:
        screen_df = None
        event_df = lf.select(event_loop_need_columns).collect()
    result = rs_event_backtest(
        capital=parameters["capital"],
        commission=parameters["commission"],
        board_lot=board_lot,
        df=event_df,
        rebalance_days_df=rebalance_days,
        record_history=record_history,
        stop_type=parameters["stop_method"]["stop_type"],
        hold_range_type=parameters["hold_method"]["hold_range_type"],
        hold_range_start=parameters["hold_method"]["hold_range_value"]["start"],
        hold_range_end=parameters["hold_method"]["hold_range_value"]["end"],
        hold_weight_type=parameters["hold_method"]["hold_weight"],
        max_hold_weight=parameters["hold_method"]["max_hold_weight"],
        hold_threshold=parameters["hold_method"]["hold_threshold"],
        hold_weight_factor=None,
    )
    if return_screen_df:
        return result[0], result[1], screen_df
    return result[0], result[1]


def _event_backtest(
    lf: pl.LazyFrame,
    parameters: BacktestParameters,
    builtin_factors: list[str],
    BuiltinScopes: Type[BuiltinScope],
    ProcessMissingDataMethods: Type[ProcessMissingDataMethod],
    NeutralizationMethods: Type[NeutralizationMethod],
    board_lot: int,
    filter_exprs: list[pl.Expr] = [],
    record_history: bool = True,
) -> tuple[pl.DataFrame, dict[str, HistoryPosition]]:
    """事件驱动回测"""
    equity_list = []
    turnover_list = []
    commission_list = []
    hold_securities_list = []

    # 交易日范围series
    trade_dates_df = (
        lf.filter(
            pl.col("trade_date").is_between(
                lower_bound=datetime.strptime(
                    parameters["date_range"]["start_date"], "%Y-%m-%d"
                ),
                upper_bound=datetime.strptime(
                    parameters["date_range"]["end_date"], "%Y-%m-%d"
                ),
            )
        )
        .select("trade_date")
        .unique()
        .sort("trade_date")
        .set_sorted("trade_date")  # 将列标记为已排序,加快未来的操作
    ).collect()

    filter_exprs.extend([pl.col("suspension"), pl.col("last_trade_day")])
    lf = (
        lf.pipe(
            planning_screen_lf,
            parameters=parameters,
            builtin_factors=builtin_factors,
            filter_exprs=filter_exprs,
            BuiltinScopes=BuiltinScopes,
            ProcessMissingDataMethods=ProcessMissingDataMethods,
            NeutralizationMethods=NeutralizationMethods,
            filter_fn=lambda x: x.filter(
                pl.col("trade_date").is_between(
                    lower_bound=datetime.strptime(
                        parameters["date_range"]["start_date"], "%Y-%m-%d"
                    ),
                    upper_bound=datetime.strptime(
                        parameters["date_range"]["end_date"], "%Y-%m-%d"
                    ),
                )
            ),
        ).sort(["trade_date", "rank"], descending=False)  # 同交易日内按rank排序
    )
    rebalance_days = calc_rebalance_days(
        frame=trade_dates_df.with_row_index(), parameters=parameters
    )
    # 初始化持仓表
    code_categories = lf.collect_schema().get("code")
    p = Portfolio(
        code_categories=code_categories,  # type: ignore
        logger=logger.event_backtest,
        balance=float(parameters["capital"]),
        commission=float(parameters["commission"]),
    )
    # logger.event_backtest.info(f"初始资金: {parameters['capital']}, 手续费率: {parameters['commission']}")
    # 事件循环所需列
    event_loop_need_columns = [
        "code",
        "name",
        "trade_date",
        "close",
        "stk_div",
        "cash_div",
        "rank",
        "limit",
        "suspension",
    ]
    # 添加持仓权重所需列
    if parameters["hold_method"]["hold_weight"] == HoldWeight.AMOUNT:
        event_loop_need_columns.append("amount")  # 添加成交量列
    # 添加止盈/止损所需列
    if parameters["stop_method"]["stop_type"] != StopType.NOTSTOP:
        _stop_trade_price_expr = (
            pl.col("pre_close") * (1 + parameters["stop_method"]["stop_value"])
        ).round(3)
        if parameters["stop_method"]["stop_type"] == StopType.TAKEPROFIT:
            # 若开盘价高于止盈价格，则止盈价格为开盘价
            lf = lf.with_columns(
                pl.when(pl.col("high") > _stop_trade_price_expr)
                .then(True)
                .otherwise(False)
                .alias("_stop_trade_flag"),
                pl.max_horizontal(_stop_trade_price_expr, pl.col("open")).alias(
                    "_stop_trade_price"
                ),
            )
        elif parameters["stop_method"]["stop_type"] == StopType.STOPLOSS:
            # 若开盘价低于止损价格，则止损价格为开盘价
            lf = lf.with_columns(
                pl.when(pl.col("low") < _stop_trade_price_expr)
                .then(True)
                .otherwise(False)
                .alias("_stop_trade_flag"),
                pl.min_horizontal(_stop_trade_price_expr, pl.col("open")).alias(
                    "_stop_trade_price"
                ),
            )
        event_loop_need_columns.extend(
            ["_stop_trade_price", "_stop_trade_flag"]
        )  # 添加止盈/止损价格列
    # logger.event_backtest.debug(f"事件循环所需列: {event_loop_need_columns}")
    # 循环
    loop_iter = (
        lf.select(event_loop_need_columns)
        .collect()
        .partition_by(
            "trade_date", maintain_order=True, include_key=False, as_dict=True
        )
        .items()
    )
    for (trade_date,), df in loop_iter:
        trade_date: date
        df: pl.DataFrame

        # 初始化计数器
        buy_amount: float = 0.0  # 本周买入金额
        sell_amount: float = 0.0  # 本周卖出金额
        buy_commission: float = 0.0  # 本周买入手续费
        sell_commission: float = 0.0  # 本周卖出手续费
        # 记录本周期开始时equity
        start_equity: float = p.equity
        # 记录本周期开始时持仓标的
        hold_securities = p.position.get_column("code")

        # 处理分红的标的
        div_df = (
            df.lazy()
            .filter(pl.col("cash_div").is_not_null() | pl.col("stk_div").is_not_null())
            .join(p.position.lazy().select("code", "volume"), on="code", how="inner")
            .select(
                pl.col("code"),
                pl.col("volume"),
                pl.col("cash_div").fill_null(0.0).alias("cash_div"),
                pl.col("stk_div").fill_null(0.0).alias("stk_div"),
            )
            .collect()
        )

        if div_df.height:
            # logger_daily.debug(f"{trade_date} 处理分红: {div_df.height}个标的")
            # 处理现金分红（扣除20%所得税）
            p.balance += (
                div_df.select(pl.col("cash_div") * pl.col("volume") / 10).sum().item()
                * 0.8
            )

            # 处理股票分红
            p.position = p.position.update(
                div_df.select(
                    pl.col("code"),
                    (
                        pl.col("volume") + (pl.col("stk_div") * pl.col("volume") / 10)
                    ).alias("volume"),
                ),
                on="code",
                how="left",
            )

            if record_history:
                # 记录分红标的
                p.history_dividend[trade_date.strftime("%Y-%m-%d")] = div_df

        # 止盈/止损模块
        if parameters["stop_method"]["stop_type"] != StopType.NOTSTOP:
            # 触发止盈/止损标的 和 持仓中的标的 的交集
            stop_trade_codes = df.filter(pl.col("_stop_trade_flag"))["code"]
            if_trigger_stop_trade = stop_trade_codes.is_in(p.position["code"]).any()
            # 若持仓中有标的触发止盈/止损
            if if_trigger_stop_trade:
                # logger_daily.debug(f"{trade_date} 触发止{'盈' if strategy_config['stop_method']['stop_type'] == StopType.TAKEPROFIT else '损'}: {stop_trade_codes.to_list()}")
                # 刷新持仓表中止盈/止损价格
                p.position = p.position.update(
                    df.filter(pl.col("_stop_trade_flag")).select(
                        pl.col("code"),
                        pl.col("_stop_trade_price").alias("last_price"),
                        pl.col("suspension"),
                    ),
                    on="code",
                    how="left",
                )
                # 用持仓中的标的生成target表
                _stop_trade_target_df = p.position.select(
                    pl.col("code"),
                    pl.col("name"),
                    pl.col("last_price"),
                    pl.col("volume"),
                    pl.when(pl.col("code").is_in(stop_trade_codes))
                    .then(pl.lit(0.0))
                    .otherwise(pl.col("volume"))
                    .alias("target_volume"),
                )
                stop_trade_target_lf = _stop_trade_target_df.select(
                    pl.exclude("volume")
                ).lazy()
                _buy_amount, _sell_amount, _buy_commission, _sell_commission = (
                    p.fill_target(stop_trade_target_lf)
                )

                buy_amount += _buy_amount  # 叠加买入金额
                sell_amount += _sell_amount  # 叠加卖出金额
                buy_commission += _buy_commission  # 叠加买入手续费
                sell_commission += _sell_commission  # 叠加卖出手续费

        # 刷新持仓表中收盘价格并更新市值
        p.position = p.position.update(
            df.select(
                pl.col("code"),
                pl.col("name"),
                pl.col("close").alias("last_price"),
                pl.col("suspension"),
            ),
            on="code",
            how="left",
        ).with_columns(
            (pl.col("last_price") * pl.col("volume")).round(2).alias("market_value")
        )

        # 判断本周期是否需要调仓
        if_rebalance = (
            rebalance_days.filter(pl.col("trade_date") == trade_date)
            .select("if_rebalance")
            .item()
        )
        if if_rebalance:
            # logger_daily.debug(f"{trade_date} 执行调仓")
            # 根据阈值修改rank
            if parameters["hold_method"]["hold_threshold"]:
                df = planning_mod_rank(
                    frame=df,
                    threshold=parameters["hold_method"]["hold_threshold"],
                    hold_range_type=parameters["hold_method"]["hold_range_type"],
                    position_lf=p.position.lazy(),
                ).collect()

            # 保留持仓范围内数据(初始化交易表)
            hold_range = parameters["hold_method"]["hold_range_value"]

            target_df = (
                df.lazy()
                .join(p.position.lazy().select("code", "volume"), on="code", how="left")
                .filter(
                    # 选出排名范围内的标的
                    pl.col("rank").is_between(
                        lower_bound=hold_range["start"],
                        upper_bound=hold_range["end"],
                        closed="both",
                    ),
                    # 排除涨停且没有持仓的标的
                    ~((pl.col("limit") == 1) & (pl.col("volume") == 0)),
                )
                .select(pl.exclude("volume"))
                .collect()
            )

            # 进行调仓
            # if target_df.height:
            # logger_daily.debug(f"目标持仓数量: {target_df.height}")
            # 计算目标持仓权重
            if parameters["hold_method"]["hold_weight"] == HoldWeight.AVG:
                _weight_expr = pl.lit(1) / target_df.height  # 等权重
            elif parameters["hold_method"]["hold_weight"] == HoldWeight.CLOSE:
                _weight_expr = pl.col("close") / pl.col("close").sum()  # 按价格加权
            elif parameters["hold_method"]["hold_weight"] == HoldWeight.AMOUNT:
                _weight_expr = pl.col("amount") / pl.col("amount").sum()  # 按成交额加权
            elif parameters["hold_method"]["hold_weight"] == HoldWeight.NOTHOLD:
                _weight_expr = pl.lit(0.0)
            else:
                raise ValueError(
                    f"Invalid hold weight: {parameters['hold_method']['hold_weight']}"
                )

            # 计算目标持仓数量(预扣除手续费后再向下取整)
            target_lf = (
                target_df.lazy()
                .select(
                    pl.col("code"),
                    pl.col("name"),
                    pl.col("close"),
                    (
                        (
                            _weight_expr.clip(
                                upper_bound=parameters["hold_method"]["max_hold_weight"]
                            )
                            * p.available_equity
                            * (1 - p.commission) ** 2
                            // (pl.col("close") * board_lot)
                        )
                        * board_lot
                    ).alias("target_volume"),
                )
                .rename({"close": "last_price"})
            )
            _buy_amount, _sell_amount, _buy_commission, _sell_commission = (
                p.fill_target(target_lf)
            )

            buy_amount += _buy_amount  # 叠加买入金额
            sell_amount += _sell_amount  # 叠加卖出金额
            buy_commission += _buy_commission  # 叠加买入手续费
            sell_commission += _sell_commission  # 叠加卖出手续费

        # 非调仓日 且触发止盈/止损 买回标的
        if (
            not if_rebalance
            and parameters["stop_method"]["stop_type"] != StopType.NOTSTOP
            and if_trigger_stop_trade
        ):
            # logger_daily.info(f"{trade_date} 非调仓日买回止{'盈' if strategy_config['stop_method']['stop_type'] == StopType.TAKEPROFIT else '损'}标的")
            # 买回止盈/止损标的
            stop_trade_target_lf = (
                _stop_trade_target_df.select(pl.exclude("target_volume"))
                .rename({"volume": "target_volume"})
                .update(
                    df.select(
                        pl.col("code"),
                        pl.col("close").alias("last_price"),  # 用收盘价填充
                    ),
                    on="code",
                    how="left",
                )
            ).lazy()
            _buy_amount, _sell_amount, _buy_commission, _sell_commission = (
                p.fill_target(stop_trade_target_lf)
            )

            buy_amount += _buy_amount  # 叠加买入金额
            sell_amount += _sell_amount  # 叠加卖出金额
            buy_commission += _buy_commission  # 叠加买入手续费
            sell_commission += _sell_commission  # 叠加卖出手续费

        if record_history:
            p.update_history_position(trade_date=trade_date.strftime("%Y-%m-%d"))

        # 保存equity, turnover, commission, hold_securities
        turnover_list.append(
            (buy_amount + sell_amount) / 2 / start_equity if start_equity else EPSILON
        )
        equity_list.append(p.equity)
        commission_list.append(buy_commission + sell_commission)
        hold_securities_list.append(hold_securities)

    # 生成res
    res = (
        trade_dates_df.lazy()
        .with_columns(
            pl.Series(name="equity", values=equity_list),
            pl.Series(name="turnover", values=turnover_list),
            pl.Series(name="commission", values=commission_list),
            pl.Series(name="hold_securities", values=hold_securities_list),
        )
        .with_columns(
            pl.col("equity")
            .pct_change()
            .fill_null(equity_list[0] / parameters["capital"] - 1)
            .alias("time_return"),  # 计算每周期回报率
        )
        .collect()
    )

    # logger.event_backtest.info(f"回测完成，最终净值: {equity_list[-1] / parameters['capital']:.4f}")
    # logger.event_backtest_daily.info(f"回测完成，最终净值: {equity_list[-1] / parameters['capital']:.4f}")
    return res, p.history_position


def cb_event_backtest(
    lf: pl.LazyFrame,
    parameters: CBBacktestParameters,
    record_history: bool = True,
    # Issue #1：透传给 _rs_event_backtest，控制是否返回全列截面数据
    return_screen_df: bool = False,
) -> tuple[pl.DataFrame, pl.DataFrame] | tuple[pl.DataFrame, pl.DataFrame, pl.DataFrame]:
    """转债事件驱动回测"""
    filter_exprs = gen_cb_filter_exprs(parameters)
    return _rs_event_backtest(
        lf=lf,
        parameters=parameters,
        builtin_factors=CB_BUILTIN_FACTORS,
        BuiltinScopes=CBBuiltinScope,
        ProcessMissingDataMethods=CBProcessMissingDataMethod,
        NeutralizationMethods=CBNeutralizationMethod,
        board_lot=10,
        filter_exprs=filter_exprs,
        record_history=record_history,
        return_screen_df=return_screen_df,
    )


def stk_event_backtest(
    lf: pl.LazyFrame,
    parameters: STKBacktestParameters,
    record_history: bool = True,
    # Issue #1：透传给 _rs_event_backtest，控制是否返回全列截面数据
    return_screen_df: bool = False,
) -> tuple[pl.DataFrame, pl.DataFrame] | tuple[pl.DataFrame, pl.DataFrame, pl.DataFrame]:
    """股票事件驱动回测"""
    filter_exprs = gen_stk_filter_exprs(parameters)
    return _rs_event_backtest(
        lf=lf,
        parameters=parameters,
        builtin_factors=STK_BUILTIN_FACTORS,
        BuiltinScopes=STKBuiltinScope,
        ProcessMissingDataMethods=STKProcessMissingDataMethod,
        NeutralizationMethods=STKNeutralizationMethod,
        board_lot=100,
        filter_exprs=filter_exprs,
        record_history=record_history,
        return_screen_df=return_screen_df,
    )


def etf_event_backtest(
    lf: pl.LazyFrame,
    parameters: ETFBacktestParameters,
    record_history: bool = True,
    # Issue #1：透传给 _rs_event_backtest，控制是否返回全列截面数据
    return_screen_df: bool = False,
) -> tuple[pl.DataFrame, pl.DataFrame] | tuple[pl.DataFrame, pl.DataFrame, pl.DataFrame]:
    """基金事件驱动回测"""
    filter_exprs = gen_etf_filter_exprs(parameters)
    return _rs_event_backtest(
        lf=lf,
        parameters=parameters,
        builtin_factors=ETF_BUILTIN_FACTORS,
        BuiltinScopes=ETFBuiltinScope,
        ProcessMissingDataMethods=ETFProcessMissingDataMethod,
        NeutralizationMethods=ETFNeutralizationMethod,
        board_lot=100,
        filter_exprs=filter_exprs,
        record_history=record_history,
        return_screen_df=return_screen_df,
    )
