from io import BytesIO

import polars as pl

from ldnode.loader.files import ETF_DATA
from ldnode.calculator.operators import etf_binning_backtest
from ldnode.calculator.plans.polars import maybe_resolve_joins
from ldnode.calculator.schemas import ETF_BUILTIN_FACTORS
from .models import ETFBinningParameters, BinningResponse, TaskResponse
from .utils import wrap_handler, response


@wrap_handler("基金分箱回测")
def handle_etf_binning(parameters: ETFBinningParameters) -> TaskResponse[BinningResponse]:
    """ 响应基金分箱回测任务 """
    bins: int = parameters.get('bins', 5) # 分箱数
    # 读取历史数据
    lf = (
        ETF_DATA.output_lazyframe()
        .filter(~pl.col('suspension'))
    )
    lf, _builtin_factors = maybe_resolve_joins(lf, 3, parameters, ETF_BUILTIN_FACTORS)
    # 回测
    return_df, corr_df, ic_mean, ir = etf_binning_backtest(lf, parameters['strategy_config'], bins, builtin_factors=_builtin_factors)
    # 将结果转成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(),
    })
