from typing import Optional

import polars as pl
from ldnode.calculator.schemas import (
    ExcludeConditionGroup,
    ExcludeFactor,
    ExcludeOperator,
)


def _collect_pct_factors(filter_factors: Optional[ExcludeConditionGroup]) -> list[str]:
    if not filter_factors:
        return []
    factors = ExcludeConditionGroup.flatten(filter_factors)
    return list(
        {
            factor["factor"]
            for factor in factors
            if "_pct" in factor["operator"]
        }
    )


def gen_exclude_expr(condition_group: ExcludeConditionGroup) -> pl.Expr:
    """生成排除公式"""
    exprs: list[pl.Expr] = []
    for condition in condition_group["conditions"]:
        if condition.get("relation"):
            # 处理条件组
            expr = gen_exclude_expr(condition)
            exprs.append(expr)
        else:
            # 处理条件
            if condition["operator"] in ExcludeOperator:
                expr = ExcludeOperator(condition["operator"]).compare(
                    condition["factor"], condition["value"]
                )
                exprs.append(expr)
    return (
        pl.Expr.or_(*exprs)
        if condition_group["relation"] == "OR"
        else pl.Expr.and_(*exprs)
    )


def planning_exclude_factors(
    frame: pl.LazyFrame | pl.DataFrame,
    filter_factors: Optional[ExcludeConditionGroup],
    filter_exprs: Optional[list[pl.Expr]],
    keep_exclude: bool = True,
) -> pl.LazyFrame:
    """
    规划排除因子计算
    排除因子的 _pct 在入池后阶段单独计算
    """
    lf = frame.lazy() if isinstance(frame, pl.DataFrame) else frame
    pct_factors = _collect_pct_factors(filter_factors)
    if pct_factors:
        lf = lf.with_columns([
            pl.col(factor)
            .rank(method="average", descending=False)
            .over("trade_date")
            .truediv(
                pl.col(factor)
                .count()
                .over("trade_date")
                .sub(pl.col(factor).null_count().over("trade_date"))
            )
            .alias(f"{factor}_pct")
            for factor in pct_factors
        ])
    # 排除因子条件
    if filter_factors:
        filter_exprs.append(gen_exclude_expr(filter_factors))
    if keep_exclude:
        lf = lf.with_columns(
            # 若filter_exprs不为空则计算_exclude列,否则直接将False赋值给_exclude列
            pl.when(pl.Expr.or_(*filter_exprs).fill_null(False))
            .then(True)
            .otherwise(pl.lit(False))
            .alias("_exclude")
            if filter_exprs
            else pl.lit(False).alias("_exclude")
        )
    elif filter_exprs:
        lf = lf.filter(map(lambda x: ~x, filter_exprs))
    return lf
