import logging
from typing import Type

import polars as pl
import polars_ds as pds

from ldnode.calculator.schemas import (
    RankFactor,
    NeutralizationMethod,
)

logger = logging.getLogger(__name__)


def _planning_discrete_process(
    df: pl.DataFrame,
    rank_factors: list[RankFactor[str, str]],
    NeutralizationMethods: Type[NeutralizationMethod],
) -> pl.DataFrame:
    """ 规划离散变量处理（在已物化的 DataFrame 上操作，避免触发额外的 I/O） """
    neutralization_factor_set = set()
    for rank_factor in rank_factors:
        if neutralization_factors := rank_factor.get('neutralization_factor'):
            neutralization_factor_set.update([
                factor for factor in neutralization_factors
                if factor in NeutralizationMethods \
                    and NeutralizationMethods(factor).is_discrete
            ])
    if neutralization_factor_set:
        # 离散化后会删除原来的因子列, 先暂存离散化后重新添加
        df = (
            df.with_columns([pl.col(factor).alias(f'{factor}_bak') for factor in neutralization_factor_set])
            .to_dummies(columns=list(neutralization_factor_set))
            .with_columns([pl.col(f'{factor}_bak').alias(factor) for factor in neutralization_factor_set])
            .drop([f'{factor}_bak' for factor in neutralization_factor_set])
        )
        # 去除离散化展开后不需要的列
        schema_names = df.columns
        cols_to_drop = [f'{factor}_' for factor in neutralization_factor_set]
        cols_to_drop.extend([f'{factor}_null' for factor in neutralization_factor_set])
        if cols_to_drop := list(filter(lambda x: x in schema_names, cols_to_drop)):
            df = df.drop(cols_to_drop)
    return df


def _gen_neutralization_expr(
    df: pl.DataFrame,
    rank_factor: RankFactor[str, str],
    NeutralizationMethods: Type[NeutralizationMethod],
) -> pl.Expr | None:
    """ 为单个 rank_factor 生成中性化表达式 """
    column_exprs: list[pl.Expr] = []
    schema_names = df.columns
    if neutralization_factors := rank_factor.get('neutralization_factor'):
        feat_idx = 0
        for factor in neutralization_factors:
            if factor in NeutralizationMethods:
                for expr in NeutralizationMethods(factor).process(schema_names):
                    # 用唯一 alias 避免与 target 列同名导致 pds.lin_reg 内部 DataFrame 重列名错误
                    column_exprs.append(expr.alias(f'_neut_feat_{feat_idx}'))
                    feat_idx += 1
    if not column_exprs:
        return None
    return (
        pds.lin_reg(
            *column_exprs,
            target=pl.col(rank_factor['factor']),
            add_bias=True,
            return_pred=True,
            null_policy='skip'
        )
        .over('trade_date')
        .struct.field('resid')
        .alias(f"{rank_factor['factor']}_neutralization")
    )


def planning_neutralization_process(
    frame: pl.LazyFrame | pl.DataFrame,
    rank_factors: list[RankFactor[str, str]],
    NeutralizationMethods: Type[NeutralizationMethod],
) -> pl.LazyFrame:
    """ 规划中性化处理 """
    lf = frame.lazy() if isinstance(frame, pl.DataFrame) else frame
    # 检查是否有因子需要中性化处理
    if all(not rank_factor.get('neutralization_factor') for rank_factor in rank_factors):
        return lf
    # 在入口处统一物化一次，后续所有操作均在内存中完成，避免中途 .collect() 重复触发 I/O
    df = lf.collect()
    # 规划离散化处理
    df = _planning_discrete_process(df, rank_factors, NeutralizationMethods)
    # 分离停牌和非停牌数据
    suspension_df = df.filter(pl.col('suspension'))
    df = df.filter(~pl.col('suspension'))
    # 逐因子执行中性化，某因子 lin_reg 因数据不足失败时回退到原始值
    # pds.lin_reg(null_policy='skip') 在有效（非NULL）数据行数 <= 特征数时
    # 会报 #Data < #features，可转债等小市场经过滤+NULL跳过后易触发
    for rank_factor in rank_factors:
        expr = _gen_neutralization_expr(df, rank_factor, NeutralizationMethods)
        if expr is None:
            continue
        try:
            df = df.with_columns(expr)
        except pl.exceptions.ComputeError as e:
            if '#Data < #features' in str(e):
                logger.warning(
                    "因子 %s 中性化因数据不足跳过（回退到原始值）: %s",
                    rank_factor['factor'], e
                )
                df = df.with_columns(
                    pl.col(rank_factor['factor']).alias(f"{rank_factor['factor']}_neutralization")
                )
            else:
                raise
    # 获取非停牌数据的schema作为目标schema
    target_schema = df.schema
    # 直接使用match_to_schema让停牌数据匹配目标schema
    suspension_df = suspension_df.lazy().match_to_schema(
        schema=target_schema,
        missing_columns='insert',  # 自动插入缺失列（包括中性化列）
        extra_columns='ignore'     # 忽略多余列
    ).collect()
    return pl.concat([df, suspension_df]).lazy().sort('trade_date')
