import polars as pl

from ldnode import logger
from ldnode.calculator.schemas import HoldMethod, HoldWeight


def require_hold_weight_factor(hold_method: HoldMethod) -> str:
    """FACTOR 权重必须显式指定因子列。"""
    factor = hold_method.get("hold_weight_factor")
    if hold_method["hold_weight"] == HoldWeight.FACTOR and not factor:
        logger.calculator.error(
            "[FactorWeight] 权重因子缺失 hold_weight=%s",
            hold_method["hold_weight"],
        )
        raise ValueError("FACTOR 持仓权重必须提供 hold_weight_factor")
    return factor or ""


def factor_raw_weight_expr(hold_method: HoldMethod) -> pl.Expr:
    factor = require_hold_weight_factor(hold_method)
    logger.calculator.info("[FactorWeight] 进入 FACTOR 权重路径 factor=%s", factor)
    return pl.col(factor)


def factor_weight_ratio_expr(hold_method: HoldMethod, include_expr: pl.Expr) -> pl.Expr:
    factor = require_hold_weight_factor(hold_method)
    raw_weight = (
        pl.when(include_expr & (pl.col(factor) > 0))
        .then(pl.col(factor))
        .otherwise(0.0)
    )
    return raw_weight.truediv(raw_weight.sum().over("trade_date"))


def log_factor_weight_stats(frame: pl.DataFrame, task_name: str, factor: str) -> None:
    if frame.is_empty():
        logger.calculator.warning(
            "[FactorWeight] 权重统计为空 task=%s factor=%s", task_name, factor
        )
        return
    stats = (
        frame.lazy()
        .group_by("trade_date")
        .agg(
            pl.len().alias("candidate_count"),
            pl.col(factor).is_null().sum().alias("null_count"),
            (pl.col(factor) < 0).sum().alias("negative_count"),
            (pl.col(factor) == 0).sum().alias("zero_count"),
            (pl.col(factor) > 0).sum().alias("valid_count"),
            pl.when(pl.col(factor) > 0)
            .then(pl.col(factor))
            .otherwise(0.0)
            .sum()
            .alias("weight_sum"),
        )
        .collect()
    )
    logger.calculator.info(
        "[FactorWeight] 权重统计 task=%s factor=%s stats=%s",
        task_name,
        factor,
        stats.to_dicts(),
    )
    invalid_dates = stats.filter(pl.col("valid_count") == 0)
    null_dates = stats.filter(pl.col("null_count") > 0)
    if null_dates.height:
        logger.calculator.error(
            "[FactorWeight] 存在空值 task=%s factor=%s dates=%s",
            task_name,
            factor,
            null_dates.select("trade_date").to_series().to_list(),
        )
        raise ValueError(f"FACTOR 权重因子存在空值: {factor}")
    negative_dates = stats.filter(pl.col("negative_count") > 0)
    if negative_dates.height:
        logger.calculator.error(
            "[FactorWeight] 存在负值 task=%s factor=%s dates=%s",
            task_name,
            factor,
            negative_dates.select("trade_date").to_series().to_list(),
        )
        raise ValueError(f"FACTOR 权重因子存在负值: {factor}")
    if invalid_dates.height:
        logger.calculator.error(
            "[FactorWeight] 全部无效 task=%s factor=%s dates=%s",
            task_name,
            factor,
            invalid_dates.select("trade_date").to_series().to_list(),
        )
        raise ValueError(f"FACTOR 权重因子全部无效: {factor}")
