from datetime import date
from typing import Optional

import polars as pl
import polars.selectors as cs

from ldnode.calculator.plans.polars import (
    planning_mod_rank,
    planning_screen_lf,
    planning_screen_weight,
)
from ldnode.calculator.schemas import (
    CB_BUILTIN_FACTORS,
    CB_SCREEN_RESULT_FIXED_COLUMNS,
    ETF_BUILTIN_FACTORS,
    ETF_SCREEN_RESULT_FIXED_COLUMNS,
    SCREEN_CB_SETTING_RELATIVE_FACTORS,
    SCREEN_COMMON_SETTING_RELATIVE_FACTORS,
    SCREEN_ETF_SETTING_RELATIVE_FACTORS,
    SCREEN_STK_SETTING_RELATIVE_FACTORS,
    STK_BUILTIN_FACTORS,
    STK_SCREEN_RESULT_FIXED_COLUMNS,
    BuiltinScope,
    CBBuiltinScope,
    CBNeutralizationMethod,
    CBProcessMissingDataMethod,
    CBScreenParameters,
    ETFBuiltinScope,
    ETFNeutralizationMethod,
    ETFProcessMissingDataMethod,
    ETFScreenParameters,
    ExcludeConditionGroup,
    NeutralizationMethod,
    ProcessMissingDataMethod,
    ScreenParameters,
    STKBuiltinScope,
    STKNeutralizationMethod,
    STKProcessMissingDataMethod,
    STKScreenParameters,
    gen_cb_filter_exprs,
    gen_etf_filter_exprs,
    gen_stk_filter_exprs,
)


def _gen_screen_result_columns(
    fixed_columns: list[str],
    parameters: ScreenParameters,
    setting_relative_factors: dict[str, list[str]],
    show_all_data: bool = False,
) -> list[str]:
    """生成返回数据列"""
    if show_all_data:
        # 显示所有列，base_columns 在前
        return (fixed_columns, cs.exclude(fixed_columns))
    # 计算要展示的自定义因子列
    customized_factors = []
    for fs in parameters.get("computed_factors", []):
        customized_factors.extend([f["factor"] for f in fs if not f["is_intermediate"]])
    # 计算快捷排除设置关联的因子列（合并通用映射 + 品种专属映射）
    merged_setting_factors = {}
    merged_setting_factors.update(SCREEN_COMMON_SETTING_RELATIVE_FACTORS)
    merged_setting_factors.update(setting_relative_factors)
    exclude_setting_factors: list[str] = []
    for setting, factors in merged_setting_factors.items():
        if parameters.get(setting):
            exclude_setting_factors.extend(factors)
    # 计算要展示的排除因子列
    if filter_factors := parameters.get("filter_factors"):
        exclude_factors = [
            f["factor"] for f in ExcludeConditionGroup.flatten(filter_factors)
        ]
    else:
        exclude_factors = []
    # 计算要展示的打分因子列
    rank_factors = [f["factor"] for f in parameters["rank_factors"]]
    # 计算要展示的打分列（命名规则与 score_factors.py 保持一致：有中性化时列名带 _neutralization）
    score_columns = [
        f"_{f['factor']}_neutralization_score" if f.get("neutralization_factor") else f"_{f['factor']}_score"
        for f in parameters["rank_factors"]
    ] + ["score"]
    # 返回列定义
    return cs.by_name(
        list(
            dict.fromkeys(
                fixed_columns
                + customized_factors
                + exclude_setting_factors
                + exclude_factors
                + rank_factors
                + score_columns
            )
        )
    )


def _planning_screen(
    screen_lf: pl.LazyFrame,
    trade_date: date,
    parameters: ScreenParameters,
    builtin_factors: list[str],
    BuiltinScopes: type[BuiltinScope],
    ProcessMissingDataMethods: type[ProcessMissingDataMethod],
    NeutralizationMethods: type[NeutralizationMethod],
    filter_exprs: list[pl.Expr],
    fixed_columns: list[str],
    setting_relative_factors: dict[str, list[str]],
    exclude_null_score_security: bool = True,
    show_all_data: bool = False,
    hold_securities: Optional[list[str]] = None,
) -> pl.LazyFrame:
    """规划选标计算"""
    # 规划选标截面
    screen_lf = screen_lf.pipe(
        planning_screen_lf,
        parameters=parameters,
        builtin_factors=builtin_factors,
        filter_exprs=filter_exprs,
        BuiltinScopes=BuiltinScopes,
        ProcessMissingDataMethods=ProcessMissingDataMethods,
        NeutralizationMethods=NeutralizationMethods,
        filter_fn=lambda x: x.filter(pl.col("trade_date") == trade_date),
    )
    # 添加权重列
    screen_lf = screen_lf.pipe(planning_screen_weight, parameters["hold_method"])
    # 若阈值不为0 且 存在持仓数据 => 规划计算修正排名
    if parameters["hold_method"]["hold_threshold"] and hold_securities:
        position_lf = pl.LazyFrame(
            {"code": hold_securities}, schema={"code": pl.String}
        ).with_columns(pl.lit(1).alias("volume"))
        screen_lf = screen_lf.pipe(
            planning_mod_rank,
            parameters["hold_method"]["hold_threshold"],
            parameters["hold_method"]["hold_range_type"],
            position_lf,
        )
    # 排除得分null的证券
    if exclude_null_score_security:
        screen_lf = screen_lf.filter(pl.col("score").is_not_null())
    screen_lf = (
        screen_lf
        # 按照rank排序
        .sort("rank", descending=False, nulls_last=True)
        # 添加选中列
        .with_columns(
            pl.col("rank")
            .is_between(
                parameters["hold_method"]["hold_range_value"]["start"],
                parameters["hold_method"]["hold_range_value"]["end"],
            )
            .fill_null(False)
            .alias("selected"),
            pl.col("time").cast(pl.String).alias("time"),  # 转换 time 为字符串
        )
        # 生成返回数据列
        .select(
            _gen_screen_result_columns(
                fixed_columns, parameters, setting_relative_factors, show_all_data
            )
        )
        # 调整返回列顺序
        .select(
            cs.by_name([c for c in fixed_columns]),  # fixed_columns 在前
            cs.exclude(fixed_columns),  # 其他列在后
        )
    )
    return screen_lf


def cb_screen(
    screen_lf: pl.LazyFrame,
    trade_date: date,
    parameters_list: list[CBScreenParameters],
    exclude_null_score_security: bool = True,
    show_all_data: bool = False,
    hold_securities: Optional[list[str]] = None,
) -> list[pl.DataFrame]:
    """选债"""
    # 生成选债计划
    screen_plans: list[pl.LazyFrame] = [
        _planning_screen(
            screen_lf=screen_lf,
            trade_date=trade_date,
            parameters=parameters,
            builtin_factors=CB_BUILTIN_FACTORS,
            BuiltinScopes=CBBuiltinScope,
            ProcessMissingDataMethods=CBProcessMissingDataMethod,
            NeutralizationMethods=CBNeutralizationMethod,
            filter_exprs=gen_cb_filter_exprs(parameters),
            fixed_columns=CB_SCREEN_RESULT_FIXED_COLUMNS,
            setting_relative_factors=SCREEN_CB_SETTING_RELATIVE_FACTORS,
            exclude_null_score_security=exclude_null_score_security,
            show_all_data=show_all_data,
            hold_securities=hold_securities,
        )
        for parameters in parameters_list
    ]
    # 执行选债计划
    return pl.collect_all(screen_plans)


def stk_screen(
    screen_lf: pl.LazyFrame,
    trade_date: date,
    parameters_list: list[STKScreenParameters],
    exclude_null_score_security: bool = True,
    show_all_data: bool = False,
    hold_securities: Optional[list[str]] = None,
):
    """选股"""
    # 生成选股计划
    screen_plans: list[pl.LazyFrame] = [
        _planning_screen(
            screen_lf=screen_lf,
            trade_date=trade_date,
            parameters=parameters,
            builtin_factors=STK_BUILTIN_FACTORS,
            BuiltinScopes=STKBuiltinScope,
            ProcessMissingDataMethods=STKProcessMissingDataMethod,
            NeutralizationMethods=STKNeutralizationMethod,
            filter_exprs=gen_stk_filter_exprs(parameters),
            fixed_columns=STK_SCREEN_RESULT_FIXED_COLUMNS,
            setting_relative_factors=SCREEN_STK_SETTING_RELATIVE_FACTORS,
            exclude_null_score_security=exclude_null_score_security,
            show_all_data=show_all_data,
            hold_securities=hold_securities,
        )
        for parameters in parameters_list
    ]
    # 执行选股计划
    return pl.collect_all(screen_plans)


def etf_screen(
    screen_lf: pl.LazyFrame,
    trade_date: date,
    parameters_list: list[ETFScreenParameters],
    exclude_null_score_security: bool = True,
    show_all_data: bool = False,
    hold_securities: Optional[list[str]] = None,
):
    """选基"""
    # 历史数据预切片：取所有策略中最大的时间窗口，只计算一次 cutoff
    max_window = max(
        (extract_max_window(p.get("computed_factors")) for p in parameters_list),
        default=0,
    )
    if max_window > 0:
        screen_lf = slice_history_by_window(screen_lf, max_window)
    # 生成选基计划
    screen_plans: list[pl.LazyFrame] = [
        _planning_screen(
            screen_lf=screen_lf,
            trade_date=trade_date,
            parameters=parameters,
            builtin_factors=ETF_BUILTIN_FACTORS,
            BuiltinScopes=ETFBuiltinScope,
            ProcessMissingDataMethods=ETFProcessMissingDataMethod,
            NeutralizationMethods=ETFNeutralizationMethod,
            filter_exprs=gen_etf_filter_exprs(parameters),
            fixed_columns=ETF_SCREEN_RESULT_FIXED_COLUMNS,
            setting_relative_factors=SCREEN_ETF_SETTING_RELATIVE_FACTORS,
            exclude_null_score_security=exclude_null_score_security,
            show_all_data=show_all_data,
            hold_securities=hold_securities,
        )
        for parameters in parameters_list
    ]
    # 执行选基计划
    return pl.collect_all(screen_plans)
