from io import BytesIO

import polars as pl

from ldnode.loader.files import STK_DATA
from ldnode.calculator.operators import stk_binning_backtest
from .models import STKBinningParameters, BinningResponse, TaskResponse
from .utils import wrap_handler, response


@wrap_handler("股票分箱回测")
def handle_stk_binning(parameters: STKBinningParameters) -> TaskResponse[BinningResponse]:
    """ 响应股票分箱回测任务 """
    bins: int = parameters.get('bins', 5) # 分箱数
    # 读取历史数据
    lf = (
        STK_DATA.output_lazyframe()
        .filter(~pl.col('suspension'))
    )
    # 回测
    return_df, corr_df, ic_mean, ir = stk_binning_backtest(lf, parameters['strategy_config'], bins)
    # 将结果转成parquet字节流
    return_buffer = BytesIO()
    return_df.write_parquet(return_buffer)
    corr_buffer = BytesIO()
    corr_df.write_parquet(corr_buffer)
    return response(data_type="binning", data={
        "ic_mean": ic_mean,
        "ir": ir,
        "return_df": return_buffer.getvalue(),
        "corr_df": corr_buffer.getvalue(),
    })
