from datetime import datetime

import polars as pl

from ldnode.calculator.operators import evaluate, stk_event_backtest
from ldnode.calculator.plans.polars import planning_earnings_curve
from ldnode.calculator.schemas.models import EvaluateData, EvaluateDataWithTest
from ldnode.loader.files import STK_DATA

from .models import STKEvalutionParameters, TaskResponse
from .utils import response, wrap_handler


def evaluate_stk(res_df: pl.DataFrame, capital: float) -> EvaluateData:
    """股票评价指标计算"""
    # 计算收益曲线
    res_df = planning_earnings_curve(
        res_df.lazy(), capital, with_benchmark=False
    ).collect()
    # 生成评价指标
    evaluation = evaluate(res_df)
    return evaluation


@wrap_handler("股票评价指标计算")
def handle_stk_evalution(
    parameters: STKEvalutionParameters,
) -> TaskResponse[EvaluateData]:
    """响应股票评价指标计算任务"""
    # 获取历史数据
    lf = STK_DATA.output_lazyframe()
    # 回测
    res_df, _ = stk_event_backtest(
        lf=lf,
        parameters=parameters["strategy_config"],
        record_history=True,
    )
    # 计算收益曲线
    res_df = planning_earnings_curve(
        res_df.lazy(),
        parameters["strategy_config"]["capital"],
        with_benchmark=False,
    ).collect()
    # 如果存在测试集起始日期 则对结果进行切割
    if (
        parameters.get("test_begin_date")
        and parameters["test_begin_date"]
        < parameters["strategy_config"]["date_range"]["end_date"]
    ):
        test_begin_date = datetime.strptime(
            parameters.get("test_begin_date"), "%Y-%m-%d"
        ).date()
        train_df = res_df.filter(pl.col("trade_date") < test_begin_date)
        test_df = res_df.filter(pl.col("trade_date") >= test_begin_date)
        train_evaluation = evaluate_stk(
            train_df, parameters["strategy_config"]["capital"]
        )
        test_evaluation = evaluate_stk(
            test_df, parameters["strategy_config"]["capital"]
        )
        evaluation = {
            **train_evaluation,
            "test": test_evaluation,
        }
    else:
        train_evaluation = evaluate_stk(
            res_df, parameters["strategy_config"]["capital"]
        )
        evaluation: EvaluateDataWithTest = {
            **train_evaluation,
        }
    return response(data_type="evalution", data=evaluation)
