from typing import Type
from datetime import datetime

import polars as pl

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



def _binning_backtest(
    lf: pl.LazyFrame,
    parameters: BacktestParameters,
    builtin_factors: list[str],
    BuiltinScopes: Type[BuiltinScope],
    ProcessMissingDataMethods: Type[ProcessMissingDataMethod],
    NeutralizationMethods: Type[NeutralizationMethod],
    bins: int,
    filter_exprs: list[pl.Expr] = [],
) -> tuple[pl.DataFrame, pl.DataFrame, float, float]:
    """ 分箱回测 """
    # 交易日范围df
    trade_dates_lf =(
        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')
    )
    if_rebalance_df = calc_rebalance_days(trade_dates_lf.with_row_index(), parameters=parameters)
    rebalance_date_df = if_rebalance_df.filter('if_rebalance').select('trade_date')
    # 生成中间环节lf
    filter_exprs.extend([pl.col('suspension')])
    stage_lf = (
        lf
        .with_columns(
            pl.col('pct_chg').shift(-1).over('code').fill_null(0).alias('return'), # 计算下周期回报
        )
        .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,
            not_rank=True,
        )
        # 根据再平衡日进行分组
        .with_columns(
            pl.when(
                pl.col('trade_date').is_in(
                    rebalance_date_df.select(pl.implode('trade_date'))['trade_date']
                )
            ).then(pl.col('trade_date')).otherwise(None).alias('trade_date_group'),
            pl.col('score')
        )
        .with_columns(
            pl.col('trade_date_group').fill_null(strategy='forward').over('code')
            .alias('trade_date') # 向前填充日期分组
        )
        .filter(pl.col('trade_date').is_not_null()) # 排除日期分组trade_date为空的行（第一次再平衡日和第二次再平衡日之间的零散数据）
        # 计算本周期涨跌幅
        .group_by(['code', 'trade_date'], maintain_order=True).agg(
            pl.col('score').first().alias('score'),
            ((pl.col('return') + 1).product() - 1).alias('return')
        )
    )
    # 计算相关性lf
    corr_lf = stage_lf.group_by('trade_date').agg(
        pl.corr(pl.col('return'), pl.col('score'), method='spearman').alias('ic')
    ).filter(pl.col('ic').is_not_nan()).sort('trade_date')

    # 收益率曲线lf
    labels=[f'{i:02d}' for i in range(1, bins+1)]
    return_lf = (
        # 分组
        stage_lf
        .with_columns(
            pl.col('score').qcut(quantiles=bins, labels=labels, allow_duplicates=True).over(pl.col('trade_date'))
            .alias('group')
        )
        .group_by(['trade_date', 'group']).agg(
            pl.col('return').mean().alias('return'), # 计算每组每周期平均回报
            pl.len().alias('count') # 计算每组每周期标的数量
        ) 
        .sort(['trade_date', 'group'])
        .with_columns(
            ((pl.col('return') + 1).cum_prod().over('group') - 1).alias('total_return') # 将当期回报转为累计回报
        )
        .select(pl.exclude('return')) # 清理无需列
    )
    return_df, corr_df = pl.collect_all([return_lf, corr_lf])
    ic_mean: float = corr_df.select('ic').mean().item()
    ir: float = ic_mean / corr_df.select('ic').std().item()
    return return_df, corr_df, ic_mean, ir


def cb_binning_backtest(
    lf: pl.LazyFrame,
    parameters: BacktestParameters,
    bins: int,
    builtin_factors: list[str] | None = None,
) -> tuple[pl.DataFrame, pl.DataFrame, float, float]:
    """ 转债分箱回测 """
    filter_exprs = gen_cb_filter_exprs(parameters)
    return _binning_backtest(
        lf=lf,
        parameters=parameters,
        builtin_factors=builtin_factors or CB_BUILTIN_FACTORS,
        BuiltinScopes=CBBuiltinScope,
        ProcessMissingDataMethods=CBProcessMissingDataMethod,
        NeutralizationMethods=CBNeutralizationMethod,
        bins=bins,
        filter_exprs=filter_exprs,
    )


def stk_binning_backtest(
    lf: pl.LazyFrame,
    parameters: BacktestParameters,
    bins: int,
    builtin_factors: list[str] | None = None,
) -> tuple[pl.DataFrame, pl.DataFrame, float, float]:
    """ 股票分箱回测 """
    filter_exprs = gen_stk_filter_exprs(parameters)
    return _binning_backtest(
        lf=lf,
        parameters=parameters,
        builtin_factors=builtin_factors or STK_BUILTIN_FACTORS,
        BuiltinScopes=STKBuiltinScope,
        ProcessMissingDataMethods=STKProcessMissingDataMethod,
        NeutralizationMethods=STKNeutralizationMethod,
        bins=bins,
        filter_exprs=filter_exprs,
    )


def etf_binning_backtest(
    lf: pl.LazyFrame,
    parameters: BacktestParameters,
    bins: int,
    builtin_factors: list[str] | None = None,
) -> tuple[pl.DataFrame, pl.DataFrame, float, float]:
    """ 基金分箱回测 """
    filter_exprs = gen_etf_filter_exprs(parameters)
    return _binning_backtest(
        lf=lf,
        parameters=parameters,
        builtin_factors=builtin_factors or ETF_BUILTIN_FACTORS,
        BuiltinScopes=ETFBuiltinScope,
        ProcessMissingDataMethods=ETFProcessMissingDataMethod,
        NeutralizationMethods=ETFNeutralizationMethod,
        bins=bins,
        filter_exprs=filter_exprs,
    )
