from typing import Type
from datetime import datetime

import polars as pl

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


def _vector_backtest(
    lf: pl.LazyFrame,
    parameters: BacktestParameters,
    builtin_factors: list[str],
    BuiltinScopes: Type[BuiltinScope],
    ProcessMissingDataMethods: Type[ProcessMissingDataMethod],
    NeutralizationMethods: Type[NeutralizationMethod],
    filter_exprs: list[pl.Expr] = [],
) -> pl.DataFrame:
    """ 向量回测 """
    # 止盈/止损模块
    if parameters['stop_method']['stop_type'] != StopType.NOTSTOP:
        stop_value = parameters['stop_method']['stop_value']
        if parameters['stop_method']['stop_type'] == StopType.TAKEPROFIT:
            _stop_trade_expr = (
                pl.when(pl.col('high_pct_chg') >= stop_value)
                .then(True).otherwise(False)
            )
            _return_expr = (
                pl.when(pl.col('high_pct_chg') < stop_value)
                .then(pl.col('pct_chg'))
                .otherwise(
                    pl.when(pl.col('open_pct_chg') > stop_value)
                    .then(pl.col('open_pct_chg'))
                    .otherwise(stop_value)
                ).shift(-1).fill_nan(0)
            )
        elif parameters['stop_method']['stop_type'] == StopType.STOPLOSS:
            _stop_trade_expr = (
                pl.when(pl.col('low_pct_chg') <= stop_value)
                .then(True).otherwise(False)
            )
            _return_expr = (
                pl.when(pl.col('low_pct_chg') > stop_value)
                .then(pl.col('pct_chg'))
                .otherwise(
                    pl.when(pl.col('open_pct_chg') < stop_value)
                    .then(pl.col('open_pct_chg'))
                    .otherwise(stop_value)
                ).shift(-1).fill_nan(0)
            )
        _exprs = [
            _stop_trade_expr.alias('_stop_trade'),
            _return_expr.alias('return')
        ]
        lf = lf.with_columns(_exprs)
    # 交易日范围series
    trade_dates_ls =(
        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')
    )
    # 阶段暂存lf
    filter_exprs.extend([pl.col('suspension')])
    stage_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')
            )),
            keep_exclude=False,
        )
        # 选出排名内的标的
        .filter(pl.col('rank').is_between(
            lower_bound=parameters['hold_method']['hold_range_value']['start'],
            upper_bound=parameters['hold_method']['hold_range_value']['end'],
        ))
        .pipe(
            planning_weight_column,
            parameters=parameters,
            with_exclude=False,
        )
        .with_columns(pl.lit(1).alias('_signal'))
        .sort(['trade_date','rank'], descending=False) # 同交易日内按rank排序
    )
    
    pivot_columns = ['trade_date', 'code', 'weight', 'return']
    if parameters['stop_method']['stop_type'] != StopType.NOTSTOP:
        pivot_columns.extend(['_stop_trade', '_signal'])
    pivot_df = (
        stage_lf
        .select(pl.col(pivot_columns))
    ).collect()
    
    # 持仓标的表
    hold_securities = pivot_df.group_by('trade_date').agg(
        pl.col('code').alias('hold_securities')
    )
    
    # 现金表
    _cash_expr = 1 - pl.col('weight').sum()
    cash_ls = trade_dates_ls.join(
        pivot_df.group_by('trade_date', maintain_order=True, ).agg(
            _cash_expr.round(2).alias('cash')
        ).lazy(),
        on='trade_date',
        how='left'
    ).fill_null(1).select(
        pl.all().exclude('trade_date')
    )
    
    # 本周期调仓后持仓权重表
    weight_after_rebalance_lf = trade_dates_ls.join(
        pivot_df
        .pivot(
            'code',
            index='trade_date',
            values='weight'
        )
        .lazy(),
        on='trade_date',
        how='left'
    ).fill_null(.0).select(
        pl.all().exclude('trade_date')
    )

    # 回报表
    return_lf = trade_dates_ls.join(
        pivot_df
        .pivot(
            'code',
            index='trade_date',
            values='return'
        )
        .lazy(),
        on='trade_date',
        how='left'
    ).fill_null(.0).select(
        pl.all().exclude('trade_date')
    ).with_columns(
        pl.lit(0).alias('cash')
    )
    if parameters['stop_method']['stop_type'] != StopType.NOTSTOP:
        stop_trade_lf = trade_dates_ls.join(
            pivot_df
            .pivot(
                'code',
                index='trade_date',
                values='_stop_trade'
            )
            .lazy(),
            on='trade_date',
            how='left'
        ).fill_null(False).select(
            pl.all().exclude('trade_date')
        )
        signal_lf = trade_dates_ls.join(
            pivot_df
            .pivot(
                'code',
                index='trade_date',
                values='_signal'
            )
            .lazy(),
            on='trade_date',
            how='left'
        ).fill_null(0).select(
            pl.all().exclude('trade_date')
        )

    if parameters['stop_method']['stop_type'] != StopType.NOTSTOP:
        trade_dates_se, cash_se, weight_after_rebalance_df, return_df, stop_trade_df, signal_df = \
            pl.collect_all([trade_dates_ls, cash_ls, weight_after_rebalance_lf, return_lf, stop_trade_lf, signal_lf])
    else:
        trade_dates_se, cash_se, weight_after_rebalance_df, return_df = \
            pl.collect_all([trade_dates_ls, cash_ls, weight_after_rebalance_lf, return_lf])
    # 合并现金
    weight_after_rebalance_df = pl.concat(
        [
            weight_after_rebalance_df,
            cash_se
        ],
        how='horizontal'
    )

    # 下周期换仓前相对权重表
    next_relative_weight_before_rebalance_df = weight_after_rebalance_df * (return_df + 1)

    # 计算组合回报(计算换手率过程中插入计算)
    _relative_equity_expr = pl.sum_horizontal(pl.all())
    _next_return_expr = _relative_equity_expr - 1
    res = next_relative_weight_before_rebalance_df.select(
        _relative_equity_expr.alias('relative_equity'),
        _next_return_expr.alias('next_return')
    )

    # 下周期调仓前权重表
    next_weight_before_rebalance_df = next_relative_weight_before_rebalance_df / res['relative_equity'] # 归一化

    diff_weight_df = weight_after_rebalance_df - next_weight_before_rebalance_df.shift(1).fill_null(0)

    if parameters['stop_method']['stop_type'] != StopType.NOTSTOP:
        pos_df = signal_df.shift(1).fill_null(0) # 本周期持仓表
        # 本周期持仓与下周期持仓变化表
        diff_pos_df = pos_df - pos_df.shift(-1).fill_null(0)
        # 本期买卖信号与下周期买卖信号是否变化表
        if_diff_pos_df = diff_pos_df != 0
        # 找出触发止盈/止损 且 买卖信号无变化标的
        stop_trade_and_not_signal_change_df = \
            (
                stop_trade_df.cast(pl.Boolean).cast(pl.Int8)
                * if_diff_pos_df.select(~pl.all()).cast(pl.Boolean).cast(pl.Int8)
            ).with_columns(
                pl.lit(0).alias('cash') # 此处需要加上现金，否则下方计算df尺寸不一致
            )
        #计算止盈/止损权重变动
        stop_diff_weight_df = \
            (next_weight_before_rebalance_df + weight_after_rebalance_df ) \
            * stop_trade_and_not_signal_change_df # 止盈/止损标志表 为 1 的行，权重设为 0
        diff_weight_df += stop_diff_weight_df # 加入止盈/止损权重变动

    # 计算换手率
    res = pl.concat([
        res,
        diff_weight_df.select(
            pl.sum_horizontal(pl.all().abs()).alias('turnover')
        ),
        ],
        how='horizontal'
    )

    _commission_expr = pl.col('turnover') * parameters['commission']
    _time_return_expr = (pl.col('next_return').shift(1).fill_null(0) + 1) * (1 - _commission_expr) - 1
    _equity_expr = (_time_return_expr + 1).cum_prod() * parameters['capital']
    res = pl.concat([
            trade_dates_se,
            res.with_columns(
                _commission_expr.alias('commission'),
                _time_return_expr.alias('time_return'), # 扣除手续费及佣金后的回报
                _equity_expr.alias('equity')
            ),
        ],
        how='horizontal'
    ).select(
        pl.all().exclude('relative_equity', 'next_return')
    ).join(
        hold_securities,
        on='trade_date',
        how='left'
    ).with_columns(
        (pl.col('commission') * pl.col('equity')).alias('commission'),
        pl.col('hold_securities').shift(1).fill_null([]).alias('hold_securities')
    )
    return res


def cb_vector_backtest(
    lf: pl.LazyFrame,
    parameters: CBBacktestParameters,
    builtin_factors: list[str] | None = None,
) -> pl.DataFrame:
    """ 可转债向量回测 """
    filter_exprs = gen_cb_filter_exprs(parameters)
    return _vector_backtest(
        lf=lf,
        parameters=parameters,
        builtin_factors=builtin_factors or CB_BUILTIN_FACTORS,
        BuiltinScopes=CBBuiltinScope,
        ProcessMissingDataMethods=CBProcessMissingDataMethod,
        NeutralizationMethods=CBNeutralizationMethod,
        filter_exprs=filter_exprs,
    )
    

def stk_vector_backtest(
    lf: pl.LazyFrame,
    parameters: STKBacktestParameters,
    builtin_factors: list[str] | None = None,
) -> pl.DataFrame:
    """ 股票向量回测 """
    filter_exprs = gen_stk_filter_exprs(parameters)
    return _vector_backtest(
        lf=lf,
        parameters=parameters,
        builtin_factors=builtin_factors or STK_BUILTIN_FACTORS,
        BuiltinScopes=STKBuiltinScope,
        ProcessMissingDataMethods=STKProcessMissingDataMethod,
        NeutralizationMethods=STKNeutralizationMethod,
        filter_exprs=filter_exprs,
    )
    

def etf_vector_backtest(
    lf: pl.LazyFrame,
    parameters: ETFBacktestParameters,
    builtin_factors: list[str] | None = None,
) -> pl.DataFrame:
    """ 基金向量回测 """
    filter_exprs = gen_etf_filter_exprs(parameters)
    return _vector_backtest(
        lf=lf,
        parameters=parameters,
        builtin_factors=builtin_factors or ETF_BUILTIN_FACTORS,
        BuiltinScopes=ETFBuiltinScope,
        ProcessMissingDataMethods=ETFProcessMissingDataMethod,
        NeutralizationMethods=ETFNeutralizationMethod,
        filter_exprs=filter_exprs,
    )
