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 _neutralize_factor(
    lf: pl.LazyFrame,
    factor: str,
    discrete_cols: list[str],
    cont_col_names: list[str],
    date_col: str = "trade_date",
) -> pl.LazyFrame:
    """对单个因子执行中性化（全程 LazyFrame，基于 FWL 定理）

    - 离散因子：group demean，等价于 OLS + dummies 残差
    - 连续因子：pds.lin_reg OLS 残差
    - 混合：先对 factor 和连续变量同步 group demean，再对 demean 后的变量做 OLS
    """
    out_col = f"{factor}_neutralization"
    fe_cols = [date_col, *discrete_cols] if discrete_cols else []

    if fe_cols and not cont_col_names:
        # 情况 1：纯离散中性化
        return lf.with_columns(
            (pl.col(factor) - pl.col(factor).mean().over(fe_cols)).alias(out_col)
        )

    if not fe_cols and cont_col_names:
        # 情况 2：纯连续中性化
        return lf.with_columns(
            pds.lin_reg(
                *(pl.col(c) for c in cont_col_names),
                target=pl.col(factor),
                add_bias=True,
                return_pred=True,
                null_policy="skip",
            )
            .over(date_col)
            .struct.field("resid")
            .alias(out_col)
        )

    if fe_cols and cont_col_names:
        # 情况 3：离散 + 连续混合（FWL 定理）
        # factor 和所有连续变量都要对离散固定效应去均值
        y_dm = f"__{factor}_dm"
        x_dm_names = [f"__{factor}_{c}_dm" for c in cont_col_names]

        lf = lf.with_columns(
            (pl.col(factor) - pl.col(factor).mean().over(fe_cols)).alias(y_dm),
            *[
                (pl.col(c) - pl.col(c).mean().over(fe_cols)).alias(dm)
                for c, dm in zip(cont_col_names, x_dm_names)
            ],
        )
        lf = lf.with_columns(
            pds.lin_reg(
                *(pl.col(c) for c in x_dm_names),
                target=pl.col(y_dm),
                add_bias=False,
                return_pred=True,
                null_policy="skip",
            )
            .over(date_col)
            .struct.field("resid")
            .alias(out_col)
        )
        return lf.drop([y_dm, *x_dm_names])

    # 情况 4：无中性化变量
    return lf.with_columns(pl.col(factor).alias(out_col))


def planning_neutralization_process(
    frame: pl.LazyFrame | pl.DataFrame,
    rank_factors: list[RankFactor[str, str]],
    NeutralizationMethods: Type[NeutralizationMethod],
) -> pl.LazyFrame:
    """规划中性化处理（全程 LazyFrame，不做 collect）"""
    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

    # 预处理：将连续中性化因子的变换表达式（如 log）提前计算为命名列，
    # 使后续 demean 和 lin_reg 能通过列名引用
    cont_transform_map: dict[str, pl.Expr] = {}
    for rank_factor in rank_factors:
        if neutralization_factors := rank_factor.get('neutralization_factor'):
            for nf in neutralization_factors:
                if nf in NeutralizationMethods and not NeutralizationMethods(nf).is_discrete:
                    if nf not in cont_transform_map:
                        exprs = NeutralizationMethods(nf).process(None)
                        if exprs:
                            col_name = f"_neut_cont_{nf}"
                            cont_transform_map[nf] = (col_name, exprs[0])

    # 将连续变量变换计算为列（如 circ_mv.log() → _neut_cont_barra_size）
    if cont_transform_map:
        lf = lf.with_columns([
            expr.alias(col_name)
            for col_name, expr in cont_transform_map.values()
        ])

    # 逐因子执行中性化
    for rank_factor in rank_factors:
        if not (neutralization_factors := rank_factor.get('neutralization_factor')):
            continue

        # 分离离散和连续中性化因子
        discrete_cols = []
        cont_col_names = []
        for nf in neutralization_factors:
            if nf not in NeutralizationMethods:
                continue
            if NeutralizationMethods(nf).is_discrete:
                discrete_cols.append(nf)
            elif nf in cont_transform_map:
                cont_col_names.append(cont_transform_map[nf][0])

        lf = _neutralize_factor(
            lf, rank_factor['factor'], discrete_cols, cont_col_names
        )

    # 清理临时列
    temp_cols = [col_name for col_name, _ in cont_transform_map.values()]
    if temp_cols:
        lf = lf.drop(temp_cols)

    return lf
