import polars as pl

from ldnode import logger
from ldnode.loader.files import ETF_DATA
from ldnode.calculator.schemas import GlanceEvaluateData
from ldnode.calculator.operators import etf_event_backtest, glance_evaluate
from .models import ETFGlanceParameters, TaskResponse
from .utils import wrap_handler, response, Timer


@wrap_handler("基金策略收益计算")
def handle_etf_glance(parameters: ETFGlanceParameters) -> TaskResponse[GlanceEvaluateData]:
    """ 响应基金策略收益计算任务 """
    # 获取历史数据
    lf = ETF_DATA.output_lazyframe()
    # 获取最新交易日期
    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"ETF回测结束. 任务ID: {parameters['task_id']} ", logger=logger.backtest_glance):
        res_df, _ = etf_event_backtest(
            lf=lf,
            parameters=parameters['strategy_config'],
            record_history=False,
        )
        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)
