from datetime import date
from typing import Optional

import polars as pl
import polars.selectors as cs

from ldnode import logger
from ldnode.calculator.plans.polars import (
    log_factor_weight_stats,
    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,
    HoldWeight,
    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"]]
    weight_factors = []
    hold_method = parameters["hold_method"]
    if hold_method["hold_weight"] == HoldWeight.FACTOR:
        weight_factors.append(hold_method["hold_weight_factor"])
    # 计算要展示的打分列（命名规则与 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
                + weight_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())
    selected_expr = pl.col("rank").is_between(
        parameters["hold_method"]["hold_range_value"]["start"],
        parameters["hold_method"]["hold_range_value"]["end"],
    )
    if parameters["hold_method"]["hold_weight"] == HoldWeight.FACTOR:
        selected_expr = selected_expr & pl.col("_weight").gt(0)

    screen_lf = (
        screen_lf
        # 按照rank排序
        .sort("rank", descending=False, nulls_last=True)
        # 添加选中列
        .with_columns(
            selected_expr.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 _collect_screen_plans(
    screen_plans: list[pl.LazyFrame],
    parameters_list: list[ScreenParameters],
    task_name: str,
) -> list[pl.DataFrame]:
    results = pl.collect_all(screen_plans)
    for parameters, result in zip(parameters_list, results):
        hold_method = parameters["hold_method"]
        if hold_method["hold_weight"] != HoldWeight.FACTOR:
            continue
        factor = hold_method["hold_weight_factor"]
        logger.screen.info("[FactorWeight] 选标 FACTOR 完成 task=%s factor=%s", task_name, factor)
        if factor in result.columns:
            factor_stats_frame = result.select(factor)
            if "trade_date" in result.columns:
                factor_stats_frame = result.select(["trade_date", factor])
            else:
                factor_stats_frame = factor_stats_frame.with_columns(
                    pl.lit(task_name).alias("trade_date")
                )
            log_factor_weight_stats(factor_stats_frame, task_name, factor)
        selected = result.filter(pl.col("selected"))
        if selected.height == 0:
            logger.screen.error("[FactorWeight] 选标全部无效 task=%s factor=%s", task_name, factor)
            raise ValueError(f"FACTOR 权重因子全部无效: {factor}")
    return results


def _drop_cb_duplicate_amount_display_columns(results: list[pl.DataFrame]) -> list[pl.DataFrame]:
    dropped_count = sum(1 for result in results if "amt" in result.columns)
    if dropped_count:
        logger.screen.info(
            f"[CBScreen] 重复成交额展示列过滤完成 removed_column=amt result_groups={dropped_count}"
        )
    return [result.drop("amt", strict=False) for result in results]


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,
    builtin_factors: list[str] | None = None,
    fixed_columns: list[str] | None = None,
) -> list[pl.DataFrame]:
    """选债"""
    # 生成选债计划
    screen_plans: list[pl.LazyFrame] = [
        _planning_screen(
            screen_lf=screen_lf,
            trade_date=trade_date,
            parameters=parameters,
            builtin_factors=builtin_factors or CB_BUILTIN_FACTORS,
            BuiltinScopes=CBBuiltinScope,
            ProcessMissingDataMethods=CBProcessMissingDataMethod,
            NeutralizationMethods=CBNeutralizationMethod,
            filter_exprs=gen_cb_filter_exprs(parameters),
            fixed_columns=fixed_columns or 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 _drop_cb_duplicate_amount_display_columns(
        _collect_screen_plans(screen_plans, parameters_list, "cb_screen")
    )


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,
    builtin_factors: list[str] | None = None,
    fixed_columns: list[str] | None = None,
):
    """选股"""
    # 生成选股计划
    screen_plans: list[pl.LazyFrame] = [
        _planning_screen(
            screen_lf=screen_lf,
            trade_date=trade_date,
            parameters=parameters,
            builtin_factors=builtin_factors or STK_BUILTIN_FACTORS,
            BuiltinScopes=STKBuiltinScope,
            ProcessMissingDataMethods=STKProcessMissingDataMethod,
            NeutralizationMethods=STKNeutralizationMethod,
            filter_exprs=gen_stk_filter_exprs(parameters),
            fixed_columns=fixed_columns or 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 _collect_screen_plans(screen_plans, parameters_list, "stk_screen")


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,
    builtin_factors: list[str] | None = None,
    fixed_columns: list[str] | None = None,
):
    """选基"""
    # 生成选基计划
    screen_plans: list[pl.LazyFrame] = [
        _planning_screen(
            screen_lf=screen_lf,
            trade_date=trade_date,
            parameters=parameters,
            builtin_factors=builtin_factors or ETF_BUILTIN_FACTORS,
            BuiltinScopes=ETFBuiltinScope,
            ProcessMissingDataMethods=ETFProcessMissingDataMethod,
            NeutralizationMethods=ETFNeutralizationMethod,
            filter_exprs=gen_etf_filter_exprs(parameters),
            fixed_columns=fixed_columns or 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 _collect_screen_plans(screen_plans, parameters_list, "etf_screen")
