import polars as pl

from ldnode import logger
from ldnode.loader.files import CB_DATA
from ldnode.calculator.schemas import CB_BUILTIN_FACTORS, GlanceEvaluateData
from ldnode.calculator.operators import cb_event_backtest, glance_evaluate
from ldnode.calculator.plans.polars import maybe_resolve_joins
from .models import CBGlanceParameters, TaskResponse
from .utils import wrap_handler, response, Timer


@wrap_handler("转债策略收益计算")
def handle_cb_glance(parameters: CBGlanceParameters) -> TaskResponse[GlanceEvaluateData]:
    """ 响应转债策略收益计算任务 """
    # 获取历史数据
    lf = (
        CB_DATA.output_lazyframe()
        # 补丁
        .with_columns(
            pl.col('code').cast(pl.String).alias('code'),
            pl.col('limit').cast(pl.Int8).alias('limit'),
            pl.col('rating').cast(pl.String).alias('rating'),
            pl.col('yy_rating').cast(pl.String).alias('yy_rating')
        )
    )
    lf, _builtin_factors = maybe_resolve_joins(lf, 1, parameters, CB_BUILTIN_FACTORS)
    # 获取最新交易日期
    latest_trade_date = lf.select(pl.col('trade_date').max()).collect().item()
    parameters['strategy_config']['date_range']['end_date'] = str(latest_trade_date)
    # 回测
    with Timer(f"转债回测结束. 任务ID: {parameters['task_id']} ", logger=logger.backtest_glance):
        res_df, _ = cb_event_backtest(
            lf=lf,
            parameters=parameters['strategy_config'],
            record_history=False,
            builtin_factors=_builtin_factors,
        )
        res_df = res_df.with_columns(
            ((pl.col('equity') - pl.col('equity').cum_max()) / pl.col('equity').cum_max()).alias('drawdown'), # 策略回撤
        )
    # 计算收益
    with Timer(f"计算收益结束. 任务ID: {parameters['task_id']} ", logger=logger.backtest_glance):
        result: GlanceEvaluateData = glance_evaluate(
            frame=res_df,
            start_date=parameters['strategy_config']['date_range']['start_date'],
            end_date=parameters['strategy_config']['date_range']['end_date'],
            created_date=parameters['strategy_config']['created_date'],
        )
    return response(data_type="glance_evaluate", data=result)
