"""Polars 翻译器 - 将 AST 翻译为 Polars 表达式。

本模块实现 DSL 到 Polars 计算图的映射，是因子表达式执行的桥梁。
支持时序因子（ts_*）和截面因子（cs_*）的防错分组翻译。
"""

from typing import Any, Callable, Dict

import numpy as np
import polars as pl

# TA-Lib 可选依赖
try:
    import talib

    HAS_TALIB = True
except ImportError:
    HAS_TALIB = False
    talib = None

from src.factors.decorators import cross_section, element_wise, time_series
from src.factors.dsl import (
    BinaryOpNode,
    Constant,
    FunctionNode,
    Node,
    Symbol,
    UnaryOpNode,
)


class PolarsTranslator:
    """Polars 表达式翻译器。

    将纯对象的 AST 树完美映射为 Polars 的带防错分组的计算图。

    Attributes:
        handlers: 函数处理器注册表，映射 func_name 到处理函数

    Example:
        >>> from src.factors.dsl import Symbol, FunctionNode
        >>> close = Symbol("close")
        >>> expr = FunctionNode("ts_mean", close, 20)
        >>> translator = PolarsTranslator()
        >>> polars_expr = translator.translate(expr)
        >>> # 结果: pl.col("close").rolling_mean(20).over("asset")
    """

    def __init__(self) -> None:
        """初始化翻译器并注册内置函数处理器。"""
        self.handlers: Dict[str, Callable[[FunctionNode], pl.Expr]] = {}
        self._register_builtin_handlers()

    def _register_builtin_handlers(self) -> None:
        """注册内置的函数处理器。"""
        # 时序因子处理器 (ts_*)
        self.register_handler("ts_mean", self._handle_ts_mean)
        self.register_handler("ts_sum", self._handle_ts_sum)
        self.register_handler("ts_std", self._handle_ts_std)
        self.register_handler("ts_max", self._handle_ts_max)
        self.register_handler("ts_min", self._handle_ts_min)
        self.register_handler("ts_delay", self._handle_ts_delay)
        self.register_handler("ts_delta", self._handle_ts_delta)
        self.register_handler("ts_corr", self._handle_ts_corr)
        self.register_handler("ts_cov", self._handle_ts_cov)
        self.register_handler("ts_var", self._handle_ts_var)
        self.register_handler("ts_skew", self._handle_ts_skew)
        self.register_handler("ts_kurt", self._handle_ts_kurt)
        self.register_handler("ts_pct_change", self._handle_ts_pct_change)
        self.register_handler("ts_ema", self._handle_ts_ema)
        self.register_handler("ts_atr", self._handle_ts_atr)
        self.register_handler("ts_rsi", self._handle_ts_rsi)
        self.register_handler("ts_obv", self._handle_ts_obv)
        self.register_handler("ts_rank", self._handle_ts_rank)
        # 补充的时间序列因子处理器
        self.register_handler("ts_sma", self._handle_ts_sma)
        self.register_handler("ts_wma", self._handle_ts_wma)
        self.register_handler("ts_decay_linear", self._handle_ts_decay_linear)
        self.register_handler("ts_argmax", self._handle_ts_argmax)
        self.register_handler("ts_argmin", self._handle_ts_argmin)
        self.register_handler("ts_count", self._handle_ts_count)
        self.register_handler("ts_prod", self._handle_ts_prod)
        self.register_handler("ts_sumac", self._handle_ts_sumac)
        self.register_handler("max", self._handle_max)
        self.register_handler("min", self._handle_min)

        # 截面因子处理器 (cs_*)

        self.register_handler("cs_rank", self._handle_cs_rank)
        self.register_handler("cs_zscore", self._handle_cs_zscore)
        self.register_handler("cs_neutral", self._handle_cs_neutral)
        self.register_handler("cs_mean", self._handle_cs_mean)
        self.register_handler("cs_winsorize", self._handle_cs_winsorize)

        # 元素级数学函数 (element_wise)
        self.register_handler("abs", self._handle_abs)
        self.register_handler("log", self._handle_log)
        self.register_handler("exp", self._handle_exp)
        self.register_handler("sqrt", self._handle_sqrt)
        self.register_handler("sign", self._handle_sign)
        self.register_handler("cos", self._handle_cos)
        self.register_handler("sin", self._handle_sin)
        self.register_handler("atan", self._handle_atan)
        self.register_handler("log1p", self._handle_log1p)
        # 条件函数
        self.register_handler("if", self._handle_if)

    def register_handler(
        self, func_name: str, handler: Callable[[FunctionNode], pl.Expr]
    ) -> None:
        """注册自定义函数处理器。

        Args:
            func_name: 函数名称
            handler: 处理函数，接收 FunctionNode 返回 pl.Expr

        Example:
            >>> def handle_custom(node: FunctionNode) -> pl.Expr:
            ...     arg = self.translate(node.args[0])
            ...     return arg * 2
            >>> translator.register_handler("custom", handle_custom)
        """
        self.handlers[func_name] = handler

    def translate(self, node: Node) -> pl.Expr:
        """递归翻译 AST 节点为 Polars 表达式。

        Args:
            node: AST 节点（Symbol、Constant、BinaryOpNode、UnaryOpNode、FunctionNode）

        Returns:
            Polars 表达式对象

        Raises:
            TypeError: 当遇到未知的节点类型时
        """
        if isinstance(node, Symbol):
            return self._translate_symbol(node)
        elif isinstance(node, Constant):
            return self._translate_constant(node)
        elif isinstance(node, BinaryOpNode):
            return self._translate_binary_op(node)
        elif isinstance(node, UnaryOpNode):
            return self._translate_unary_op(node)
        elif isinstance(node, FunctionNode):
            return self._translate_function(node)
        else:
            raise TypeError(f"未知的节点类型: {type(node).__name__}")

    def _translate_symbol(self, node: Symbol) -> pl.Expr:
        """翻译 Symbol 节点为 pl.col() 表达式。

        Args:
            node: 符号节点

        Returns:
            pl.col(node.name) 表达式
        """
        return pl.col(node.name)

    def _translate_constant(self, node: Constant) -> pl.Expr:
        """翻译 Constant 节点为 Polars 字面量。

        Args:
            node: 常量节点

        Returns:
            pl.lit(node.value) 表达式
        """
        return pl.lit(node.value)

    def _translate_binary_op(self, node: BinaryOpNode) -> pl.Expr:
        """翻译 BinaryOpNode 为 Polars 二元运算。

        Args:
            node: 二元运算节点

        Returns:
            Polars 二元运算表达式
        """
        left = self.translate(node.left)
        right = self.translate(node.right)

        op_map = {
            "+": lambda l, r: l + r,
            "-": lambda l, r: l - r,
            "*": lambda l, r: l * r,
            # 【修复】除法处理除零，避免产生 NaN/inf 导致 EMA 永久感染
            "/": lambda l, r: pl.when(r == 0).then(None).otherwise(l / r),
            "**": lambda l, r: l.pow(r),
            "//": lambda l, r: l.floor_div(r),
            "%": lambda l, r: l % r,
            "==": lambda l, r: l.eq(r),
            "!=": lambda l, r: l.ne(r),
            "<": lambda l, r: l.lt(r),
            "<=": lambda l, r: l.le(r),
            ">": lambda l, r: l.gt(r),
            ">=": lambda l, r: l.ge(r),
            "&": lambda l, r: l & r,
            "|": lambda l, r: l | r,
        }

        if node.op not in op_map:
            raise ValueError(f"不支持的二元运算符: {node.op}")

        return op_map[node.op](left, right)

    def _translate_unary_op(self, node: UnaryOpNode) -> pl.Expr:
        """翻译 UnaryOpNode 为 Polars 一元运算。

        Args:
            node: 一元运算节点

        Returns:
            Polars 一元运算表达式
        """
        operand = self.translate(node.operand)

        op_map = {
            "+": lambda x: x,
            "-": lambda x: -x,
            "abs": lambda x: x.abs(),
        }

        if node.op not in op_map:
            raise ValueError(f"不支持的一元运算符: {node.op}")

        return op_map[node.op](operand)

    def _translate_function(self, node: FunctionNode) -> pl.Expr:
        """翻译 FunctionNode 为 Polars 函数调用。

        优先从 handlers 注册表中查找处理器，未找到则抛出错误。

        Args:
            node: 函数调用节点

        Returns:
            Polars 函数表达式

        Raises:
            ValueError: 当函数名称未注册处理器时
        """
        func_name = node.func_name

        if func_name in self.handlers:
            return self.handlers[func_name](node)
        else:
            raise ValueError(
                f"未注册的函数: {func_name}. 请使用 register_handler 注册处理器。"
            )

    # ==================== 时序因子处理器 (ts_*) ====================
    # 所有时序因子使用 @time_series 装饰器自动注入 over("ts_code") 防串表

    @time_series
    def _handle_ts_mean(self, node: FunctionNode) -> pl.Expr:
        """处理 ts_mean(close, window) -> rolling_mean(window)。"""
        if len(node.args) != 2:
            raise ValueError("ts_mean 需要 2 个参数: (expr, window)")
        expr = self.translate(node.args[0])
        window = self._extract_window(node.args[1])
        return expr.rolling_mean(window_size=window)

    @time_series
    def _handle_ts_sum(self, node: FunctionNode) -> pl.Expr:
        """处理 ts_sum(close, window) -> rolling_sum(window)。"""
        if len(node.args) != 2:
            raise ValueError("ts_sum 需要 2 个参数: (expr, window)")
        expr = self.translate(node.args[0])
        window = self._extract_window(node.args[1])
        return expr.rolling_sum(window_size=window)

    @time_series
    def _handle_ts_std(self, node: FunctionNode) -> pl.Expr:
        """处理 ts_std(close, window) -> rolling_std(window)。"""
        if len(node.args) != 2:
            raise ValueError("ts_std 需要 2 个参数: (expr, window)")
        expr = self.translate(node.args[0])
        window = self._extract_window(node.args[1])
        return expr.rolling_std(window_size=window)

    @time_series
    def _handle_ts_max(self, node: FunctionNode) -> pl.Expr:
        """处理 ts_max(close, window) -> rolling_max(window)。"""
        if len(node.args) != 2:
            raise ValueError("ts_max 需要 2 个参数: (expr, window)")
        expr = self.translate(node.args[0])
        window = self._extract_window(node.args[1])
        return expr.rolling_max(window_size=window)

    @time_series
    def _handle_ts_min(self, node: FunctionNode) -> pl.Expr:
        """处理 ts_min(close, window) -> rolling_min(window)。"""
        if len(node.args) != 2:
            raise ValueError("ts_min 需要 2 个参数: (expr, window)")
        expr = self.translate(node.args[0])
        window = self._extract_window(node.args[1])
        return expr.rolling_min(window_size=window)

    @time_series
    def _handle_ts_delay(self, node: FunctionNode) -> pl.Expr:
        """处理 ts_delay(close, n) -> shift(n)。"""
        if len(node.args) != 2:
            raise ValueError("ts_delay 需要 2 个参数: (expr, n)")
        expr = self.translate(node.args[0])
        n = self._extract_window(node.args[1])
        return expr.shift(n)

    @time_series
    def _handle_ts_delta(self, node: FunctionNode) -> pl.Expr:
        """处理 ts_delta(close, n) -> (expr - shift(n))。"""
        if len(node.args) != 2:
            raise ValueError("ts_delta 需要 2 个参数: (expr, n)")
        expr = self.translate(node.args[0])
        n = self._extract_window(node.args[1])
        return expr - expr.shift(n)

    @time_series
    def _handle_ts_corr(self, node: FunctionNode) -> pl.Expr:
        """处理 ts_corr(x, y, window) -> rolling_corr(x, y, window_size)。"""
        if len(node.args) != 3:
            raise ValueError("ts_corr 需要 3 个参数: (x, y, window)")
        x = self.translate(node.args[0])
        y = self.translate(node.args[1])
        window = self._extract_window(node.args[2])
        return pl.rolling_corr(x, y, window_size=window)

    @time_series
    def _handle_ts_cov(self, node: FunctionNode) -> pl.Expr:
        """处理 ts_cov(x, y, window) -> rolling_cov(x, y, window_size)。"""
        if len(node.args) != 3:
            raise ValueError("ts_cov 需要 3 个参数: (x, y, window)")
        x = self.translate(node.args[0])
        y = self.translate(node.args[1])
        window = self._extract_window(node.args[2])
        return pl.rolling_cov(x, y, window_size=window)

    @time_series
    def _handle_ts_var(self, node: FunctionNode) -> pl.Expr:
        """处理 ts_var(close, window) -> rolling_var(window)。"""
        if len(node.args) != 2:
            raise ValueError("ts_var 需要 2 个参数: (expr, window)")
        expr = self.translate(node.args[0])
        window = self._extract_window(node.args[1])
        return expr.rolling_var(window_size=window)

    @time_series
    def _handle_ts_skew(self, node: FunctionNode) -> pl.Expr:
        """处理 ts_skew(close, window) -> rolling_skew(window)。"""
        if len(node.args) != 2:
            raise ValueError("ts_skew 需要 2 个参数: (expr, window)")
        expr = self.translate(node.args[0])
        window = self._extract_window(node.args[1])
        return expr.rolling_skew(window_size=window)

    @time_series
    def _handle_ts_kurt(self, node: FunctionNode) -> pl.Expr:
        """处理 ts_kurt(close, window) -> rolling_kurt(window)。"""
        if len(node.args) != 2:
            raise ValueError("ts_kurt 需要 2 个参数: (expr, window)")
        expr = self.translate(node.args[0])
        window = self._extract_window(node.args[1])

        # 抛弃极慢的 rolling_map，借用 pandas 的 Cython 引擎
        def kurt_calc(s: pl.Series) -> pl.Series:
            import pandas as pd

            # pandas.rolling.kurt() 是用 Cython 编写的，速度比 pure python 快很多
            pd_series = pd.Series(s.to_numpy())
            result = pd_series.rolling(window).kurt().to_numpy()
            return pl.Series(result)

        return expr.map_batches(kurt_calc, return_dtype=pl.Float64)

    @time_series
    def _handle_ts_pct_change(self, node: FunctionNode) -> pl.Expr:
        """处理 ts_pct_change(x, n) -> (x - shift(n)) / shift(n)。"""
        if len(node.args) != 2:
            raise ValueError("ts_pct_change 需要 2 个参数: (expr, periods)")
        expr = self.translate(node.args[0])
        n = self._extract_window(node.args[1])
        shifted = expr.shift(n)
        return (expr - shifted) / shifted

    @time_series
    def _handle_ts_ema(self, node: FunctionNode) -> pl.Expr:
        """处理 ts_ema(x, window) -> ewm_mean(span=window)。"""
        if len(node.args) != 2:
            raise ValueError("ts_ema 需要 2 个参数: (expr, window)")
        expr = self.translate(node.args[0])
        window = self._extract_window(node.args[1])
        return expr.ewm_mean(span=window)

    @time_series
    def _handle_ts_atr(self, node: FunctionNode) -> pl.Expr:
        """处理 ts_atr(high, low, close, window) -> 使用 TA-Lib 计算 ATR。

        使用 map_batches 在每个分组上应用 TA-Lib ATR 函数。
        @time_series 装饰器会自动添加 .over("ts_code")
        """
        if not HAS_TALIB:
            raise ImportError("ts_atr 需要安装 TA-Lib。请运行: pip install TA-Lib")
        if len(node.args) != 4:
            raise ValueError("ts_atr 需要 4 个参数: (high, low, close, window)")

        high = self.translate(node.args[0])
        low = self.translate(node.args[1])
        close = self.translate(node.args[2])
        window = self._extract_window(node.args[3])

        # 使用 map_batches 应用 TA-Lib ATR 到整个分组
        def calc_atr(struct_series: pl.Series) -> pl.Series:
            """计算 ATR 的辅助函数。"""
            if len(struct_series) == 0:
                return pl.Series([float("nan")] * len(struct_series))

            # struct_series 包含 h, l, c 三个字段
            h = np.array(struct_series.struct.field("h").to_list(), dtype=float)
            l = np.array(struct_series.struct.field("l").to_list(), dtype=float)
            c = np.array(struct_series.struct.field("c").to_list(), dtype=float)
            result = talib.ATR(h, l, c, timeperiod=window)
            return pl.Series(result)

        return pl.struct(
            [high.alias("h"), low.alias("l"), close.alias("c")]
        ).map_batches(calc_atr, return_dtype=pl.Float64)

    @time_series
    def _handle_ts_rsi(self, node: FunctionNode) -> pl.Expr:
        """处理 ts_rsi(close, window) -> 使用 TA-Lib 计算 RSI。

        使用 map_batches 在每个分组上应用 TA-Lib RSI 函数。
        @time_series 装饰器会自动添加 .over("ts_code")
        """
        if not HAS_TALIB:
            raise ImportError("ts_rsi 需要安装 TA-Lib。请运行: pip install TA-Lib")
        if len(node.args) != 2:
            raise ValueError("ts_rsi 需要 2 个参数: (close, window)")

        close = self.translate(node.args[0])
        window = self._extract_window(node.args[1])

        # 使用 map_batches 应用 TA-Lib RSI 到整个分组
        def calc_rsi(series: pl.Series) -> pl.Series:
            """计算 RSI 的辅助函数。"""
            values = np.array(series.to_list(), dtype=float)
            result = talib.RSI(values, timeperiod=window)
            return pl.Series(result)

        return close.map_batches(calc_rsi, return_dtype=pl.Float64)

    @time_series
    def _handle_ts_obv(self, node: FunctionNode) -> pl.Expr:
        """处理 ts_obv(close, volume) -> 使用 TA-Lib 计算 OBV。

        使用 map_batches 在每个分组上应用 TA-Lib OBV 函数。
        @time_series 装饰器会自动添加 .over("ts_code")
        """
        if not HAS_TALIB:
            raise ImportError("ts_obv 需要安装 TA-Lib。请运行: pip install TA-Lib")
        if len(node.args) != 2:
            raise ValueError("ts_obv 需要 2 个参数: (close, volume)")

        close = self.translate(node.args[0])
        volume = self.translate(node.args[1])

        # 使用 map_batches 应用 TA-Lib OBV 到整个分组
        def calc_obv(struct_series: pl.Series) -> pl.Series:
            """计算 OBV 的辅助函数。"""
            if len(struct_series) == 0:
                return pl.Series([float("nan")] * len(struct_series))

            # struct_series 包含 c 和 v 两个字段
            c = np.array(struct_series.struct.field("c").to_list(), dtype=float)
            v = np.array(struct_series.struct.field("v").to_list(), dtype=float)
            result = talib.OBV(c, v)
            return pl.Series(result)

        return pl.struct([close.alias("c"), volume.alias("v")]).map_batches(
            calc_obv, return_dtype=pl.Float64
        )

    @time_series
    def _handle_ts_rank(self, node: FunctionNode) -> pl.Expr:
        """处理 ts_rank(x, window) -> 滚动排名（分位数）。"""
        if len(node.args) != 2:
            raise ValueError("ts_rank 需要 2 个参数: (x, window)")
        expr = self.translate(node.args[0])
        window = self._extract_window(node.args[1])

        def rank_calc(s: pl.Series) -> pl.Series:
            values = s.to_numpy()
            n = len(values)
            if n < window:
                return pl.Series([float("nan")] * n)

            # 核心魔法：创建零拷贝的 2D 滑动窗口视图
            # 形状为 (N - window + 1, window)
            windows = np.lib.stride_tricks.sliding_window_view(values, window)

            # 当前值即为每个窗口的最后一个元素 (N - window + 1, )
            current_vals = windows[:, -1]

            # 【终极修复】使用窗口内实际有效数据个数作为分母
            # 1. 统计小于等于当前值的个数
            less_equal = np.sum(windows <= current_vals[:, None], axis=1)
            # 2. 统计当前窗口内有效的非 NaN 数据个数
            valid_counts = np.sum(~np.isnan(windows), axis=1)

            # 3. 使用真实有效个数作为分母，避免分母陷阱
            with np.errstate(divide="ignore", invalid="ignore"):
                ranks = np.where(valid_counts > 0, less_equal / valid_counts, np.nan)

            # 【修复】如果当前值是 NaN，则排名也必须是 NaN
            ranks[np.isnan(current_vals)] = np.nan

            result = np.full(n, np.nan)
            result[window - 1 :] = ranks
            return pl.Series(result)

        return expr.map_batches(rank_calc, return_dtype=pl.Float64)

    # ==================== 补充的时间序列因子处理器 ====================

    @time_series
    def _handle_ts_sma(self, node: FunctionNode) -> pl.Expr:
        """处理 ts_sma(x, n, m) -> ewm_mean(alpha=m/n, adjust=False)。

        国泰君安 191 的 SMA 使用 N*M 平滑，等效于 alpha=m/n 的 EWM。
        """
        if len(node.args) != 3:
            raise ValueError("ts_sma 需要 3 个参数: (x, n, m)")
        expr = self.translate(node.args[0])
        n = self._extract_window(node.args[1])
        m = self._extract_window(node.args[2])
        alpha = m / n if n != 0 else 0.5
        return expr.ewm_mean(alpha=alpha, adjust=False)

    @time_series
    def _handle_ts_wma(self, node: FunctionNode) -> pl.Expr:
        """处理 ts_wma(x, window) -> 使用 numpy.convolve 实现线性加权。

        使用卷积实现高性能线性加权移动平均。
        """
        if len(node.args) != 2:
            raise ValueError("ts_wma 需要 2 个参数: (x, window)")
        expr = self.translate(node.args[0])
        window = self._extract_window(node.args[1])

        def wma_calc(s: pl.Series) -> pl.Series:
            """计算线性加权移动平均，使用 numpy 卷积优化。"""
            values = s.to_numpy()
            n = len(values)
            if n == 0 or n < window:
                # 数据为空或不足以计算完整窗口时，返回全 NaN
                return pl.Series([float("nan")] * n)

            # 线性递增权重: 1, 2, 3, ..., window
            weights = np.arange(1, window + 1, dtype=float)
            weights = weights / weights.sum()  # 归一化

            # 使用卷积计算加权平均
            result = np.full(n, np.nan)
            # valid 卷积结果需要至少 window 个有效数据点
            conv_res = np.convolve(values, weights[::-1], mode="valid")
            result[window - 1 :] = conv_res

            return pl.Series(result)

        return expr.map_batches(wma_calc, return_dtype=pl.Float64)

    @time_series
    def _handle_ts_decay_linear(self, node: FunctionNode) -> pl.Expr:
        """处理 ts_decay_linear(x, window) -> 复用 ts_wma 实现。

        DECAYLINEAR 与 WMA 等价，都使用近期最高权重的线性权重。
        直接复用 ts_wma 的实现。
        """
        # 复用 ts_wma 的实现
        return self._handle_ts_wma(node)

    @time_series
    def _handle_ts_argmax(self, node: FunctionNode) -> pl.Expr:
        """处理 ts_argmax(x, window) -> 距离最大值的交易日数。"""
        if len(node.args) != 2:
            raise ValueError("ts_argmax 需要 2 个参数: (x, window)")
        expr = self.translate(node.args[0])
        window = self._extract_window(node.args[1])

        def argmax_calc(s: pl.Series) -> pl.Series:
            values = s.to_numpy()
            n = len(values)
            if n < window:
                return pl.Series([float("nan")] * n)

            # 创建 2D 视图
            windows = np.lib.stride_tricks.sliding_window_view(values, window)

            # 沿窗口轴(axis=1)进行 C 级极速 argmax 查找
            # 注意：若有缺失值，可用 np.nanargmax 防止报错
            argmax_indices = np.nanargmax(windows, axis=1)

            # 计算距离
            distances = window - 1 - argmax_indices

            result = np.full(n, np.nan)
            result[window - 1 :] = distances
            return pl.Series(result)

        return expr.map_batches(argmax_calc, return_dtype=pl.Float64)

    @time_series
    def _handle_ts_argmin(self, node: FunctionNode) -> pl.Expr:
        """处理 ts_argmin(x, window) -> 距离最小值的交易日数。"""
        if len(node.args) != 2:
            raise ValueError("ts_argmin 需要 2 个参数: (x, window)")
        expr = self.translate(node.args[0])
        window = self._extract_window(node.args[1])

        def argmin_calc(s: pl.Series) -> pl.Series:
            values = s.to_numpy()
            n = len(values)
            if n < window:
                return pl.Series([float("nan")] * n)

            windows = np.lib.stride_tricks.sliding_window_view(values, window)
            argmin_indices = np.nanargmin(windows, axis=1)
            distances = window - 1 - argmin_indices

            result = np.full(n, np.nan)
            result[window - 1 :] = distances
            return pl.Series(result)

        return expr.map_batches(argmin_calc, return_dtype=pl.Float64)

    @time_series
    def _handle_ts_count(self, node: FunctionNode) -> pl.Expr:
        """处理 ts_count(condition, window) -> 布尔expr.cast(Int32).rolling_sum(window)。"""
        if len(node.args) != 2:
            raise ValueError("ts_count 需要 2 个参数: (condition, window)")
        expr = self.translate(node.args[0])
        window = self._extract_window(node.args[1])
        # 布尔型转换为 Int32 后 rolling_sum
        return expr.cast(pl.Int32).rolling_sum(window_size=window)

    @time_series
    def _handle_ts_prod(self, node: FunctionNode) -> pl.Expr:
        """处理 ts_prod(x, window) -> 滚动连乘积。"""
        if len(node.args) != 2:
            raise ValueError("ts_prod 需要 2 个参数: (x, window)")
        expr = self.translate(node.args[0])
        window = self._extract_window(node.args[1])

        def prod_calc(s: pl.Series) -> pl.Series:
            values = s.to_numpy()
            n = len(values)
            if n < window:
                return pl.Series([float("nan")] * n)

            windows = np.lib.stride_tricks.sliding_window_view(values, window)
            # 沿轴求积，极速实现
            prods = np.prod(windows, axis=1)

            result = np.full(n, np.nan)
            result[window - 1 :] = prods
            return pl.Series(result)

        return expr.map_batches(prod_calc, return_dtype=pl.Float64)

    @time_series
    def _handle_ts_sumac(self, node: FunctionNode) -> pl.Expr:
        """处理 ts_sumac(x) -> cum_sum()。"""
        if len(node.args) != 1:
            raise ValueError("ts_sumac 需要 1 个参数: (x)")
        expr = self.translate(node.args[0])
        return expr.cum_sum()

    @element_wise
    def _handle_max(self, node: FunctionNode) -> pl.Expr:
        """处理 max(x, y) -> 逐元素最大值。"""
        if len(node.args) != 2:
            raise ValueError("max 需要 2 个参数: (x, y)")
        x = self.translate(node.args[0])
        y = self.translate(node.args[1])
        return pl.when(x >= y).then(x).otherwise(y)

    @element_wise
    def _handle_min(self, node: FunctionNode) -> pl.Expr:
        """处理 min(x, y) -> 逐元素最小值。"""
        if len(node.args) != 2:
            raise ValueError("min 需要 2 个参数: (x, y)")
        x = self.translate(node.args[0])
        y = self.translate(node.args[1])
        return pl.when(x <= y).then(x).otherwise(y)

    @element_wise
    def _handle_if(self, node: FunctionNode) -> pl.Expr:
        """处理 if(condition, true_val, false_val) -> 条件选择。

        使用 Polars 的 when/then/otherwise 实现条件选择。
        """
        if len(node.args) != 3:
            raise ValueError("if 需要 3 个参数: (condition, true_val, false_val)")
        condition = self.translate(node.args[0])
        true_val = self.translate(node.args[1])
        false_val = self.translate(node.args[2])
        return pl.when(condition).then(true_val).otherwise(false_val)

    # ==================== 截面因子处理器 (cs_*) ====================
    # 所有截面因子使用 @cross_section 装饰器自动注入 over("trade_date") 防串表

    @cross_section
    def _handle_cs_rank(self, node: FunctionNode) -> pl.Expr:
        """处理 cs_rank(expr) -> rank()/count()。

        将排名归一化到 [0, 1] 区间。
        """
        if len(node.args) != 1:
            raise ValueError("cs_rank 需要 1 个参数: (expr)")
        expr = self.translate(node.args[0])
        return expr.rank() / expr.count()

    @cross_section
    def _handle_cs_zscore(self, node: FunctionNode) -> pl.Expr:
        """处理 cs_zscore(expr) -> (expr - mean())/std()。"""
        if len(node.args) != 1:
            raise ValueError("cs_zscore 需要 1 个参数: (expr)")
        expr = self.translate(node.args[0])
        return (expr - expr.mean()) / expr.std()

    @cross_section
    def _handle_cs_neutral(self, node: FunctionNode) -> pl.Expr:
        """处理 cs_neutral(expr, group) -> 分组中性化。"""
        if len(node.args) not in [1, 2]:
            raise ValueError("cs_neutral 需要 1-2 个参数: (expr, [group_col])")
        expr = self.translate(node.args[0])
        # 简单实现：减去截面均值（可在未来扩展为分组中性化）
        return expr - expr.mean()

    @cross_section
    def _handle_cs_mean(self, node: FunctionNode) -> pl.Expr:
        """处理 cs_mean(expr) -> 截面均值。

        计算因子在横截面上的平均值，常用于 Alpha127 等因子。
        例如：MEAN((100*(CLOSE-MAX(CLOSE,12))/(MAX(CLOSE,12)))^2) 中的 MEAN
        """
        if len(node.args) != 1:
            raise ValueError("cs_mean 需要 1 个参数: (expr)")
        expr = self.translate(node.args[0])
        return expr.mean()

    @cross_section
    def _handle_cs_winsorize(self, node: FunctionNode) -> pl.Expr:
        """处理 cs_winsorize(expr, lower, upper) -> 截面缩尾处理。"""
        if len(node.args) not in [1, 3]:
            raise ValueError("cs_winsorize 需要 1 或 3 个参数: (expr, [lower, upper])")
        expr = self.translate(node.args[0])
        if len(node.args) == 3:
            lower = self._extract_float(node.args[1])
            upper = self._extract_float(node.args[2])
        else:
            lower = 0.01
            upper = 0.99
        lower_q = expr.quantile(lower)
        upper_q = expr.quantile(upper)
        return (
            pl.when(expr < lower_q)
            .then(lower_q)
            .when(expr > upper_q)
            .then(upper_q)
            .otherwise(expr)
        )

    # ==================== 元素级数学函数 (element_wise) ====================
    # 这些函数对每个元素独立计算，不添加 over

    @element_wise
    def _handle_log(self, node: FunctionNode) -> pl.Expr:
        """处理 log(expr) -> 自然对数。"""
        if len(node.args) != 1:
            raise ValueError("log 需要 1 个参数: (expr)")
        expr = self.translate(node.args[0])
        return expr.log()

    @element_wise
    def _handle_abs(self, node: FunctionNode) -> pl.Expr:
        """处理 abs(expr) -> 绝对值。"""
        if len(node.args) != 1:
            raise ValueError("abs 需要 1 个参数: (expr)")
        expr = self.translate(node.args[0])
        return expr.abs()

    @element_wise
    def _handle_exp(self, node: FunctionNode) -> pl.Expr:
        """处理 exp(expr) -> 指数函数。"""
        if len(node.args) != 1:
            raise ValueError("exp 需要 1 个参数: (expr)")
        expr = self.translate(node.args[0])
        return expr.exp()

    @element_wise
    def _handle_sqrt(self, node: FunctionNode) -> pl.Expr:
        """处理 sqrt(expr) -> 平方根。"""
        if len(node.args) != 1:
            raise ValueError("sqrt 需要 1 个参数: (expr)")
        expr = self.translate(node.args[0])
        return expr.sqrt()

    @element_wise
    def _handle_sign(self, node: FunctionNode) -> pl.Expr:
        """处理 sign(expr) -> 符号函数。"""
        if len(node.args) != 1:
            raise ValueError("sign 需要 1 个参数: (expr)")
        expr = self.translate(node.args[0])
        return expr.sign()

    @element_wise
    def _handle_cos(self, node: FunctionNode) -> pl.Expr:
        """处理 cos(expr) -> 余弦函数。"""
        if len(node.args) != 1:
            raise ValueError("cos 需要 1 个参数: (expr)")
        expr = self.translate(node.args[0])
        return expr.cos()

    @element_wise
    def _handle_sin(self, node: FunctionNode) -> pl.Expr:
        """处理 sin(expr) -> 正弦函数。"""
        if len(node.args) != 1:
            raise ValueError("sin 需要 1 个参数: (expr)")
        expr = self.translate(node.args[0])
        return expr.sin()

    @element_wise
    def _handle_atan(self, node: FunctionNode) -> pl.Expr:
        """处理 atan(expr) -> 反正切函数。"""
        if len(node.args) != 1:
            raise ValueError("atan 需要 1 个参数: (expr)")
        expr = self.translate(node.args[0])
        return expr.arctan()

    @element_wise
    def _handle_log1p(self, node: FunctionNode) -> pl.Expr:
        """处理 log1p(expr) -> log(1+x) 函数。"""
        if len(node.args) != 1:
            raise ValueError("log1p 需要 1 个参数: (expr)")
        expr = self.translate(node.args[0])
        return expr.log1p()

    # ==================== 辅助方法 ====================

    def _extract_window(self, node: Node) -> int:
        """从节点中提取窗口大小参数。

        Args:
            node: 应该是 Constant 节点

        Returns:
            整数值

        Raises:
            ValueError: 当节点不是 Constant 或值不是整数时
        """
        if isinstance(node, Constant):
            if not isinstance(node.value, int):
                raise ValueError(
                    f"窗口参数必须是整数，得到: {type(node.value).__name__}"
                )
            return node.value
        raise ValueError(f"窗口参数必须是常量整数，得到: {type(node).__name__}")

    def _extract_float(self, node: Node) -> float:
        """从节点中提取浮点参数。

        Args:
            node: 应该是 Constant 节点

        Returns:
            float 值

        Raises:
            ValueError: 当节点不是 Constant 或值不是数值时
        """
        if isinstance(node, Constant):
            if not isinstance(node.value, (int, float)):
                raise ValueError(
                    f"数值参数必须是 int 或 float，得到: {type(node.value).__name__}"
                )
            return float(node.value)
        raise ValueError(f"数值参数必须是常量，得到: {type(node).__name__}")


def translate_to_polars(node: Node) -> pl.Expr:
    """便捷函数 - 将 AST 节点翻译为 Polars 表达式。

    Args:
        node: 表达式树的根节点

    Returns:
        Polars 表达式对象

    Example:
        >>> from src.factors.dsl import Symbol, FunctionNode
        >>> close = Symbol("close")
        >>> expr = FunctionNode("ts_mean", close, 20)
        >>> polars_expr = translate_to_polars(expr)
    """
    translator = PolarsTranslator()
    return translator.translate(node)


if __name__ == "__main__":
    # 测试用例
    from src.factors.dsl import Symbol, FunctionNode

    # 创建符号
    close = Symbol("close")
    volume = Symbol("volume")

    # 测试 1: 简单符号
    print("测试 1: Symbol")
    translator = PolarsTranslator()
    expr1 = translator.translate(close)
    print(f"  close -> {expr1}")
    assert str(expr1) == 'col("close")'

    # 测试 2: 二元运算
    print("\n测试 2: BinaryOp")
    expr2 = translator.translate(close + 10)
    print(f"  close + 10 -> {expr2}")

    # 测试 3: ts_mean
    print("\n测试 3: ts_mean")
    expr3 = translator.translate(FunctionNode("ts_mean", close, 20))
    print(f"  ts_mean(close, 20) -> {expr3}")

    # 测试 4: cs_rank
    print("\n测试 4: cs_rank")
    expr4 = translator.translate(FunctionNode("cs_rank", close / volume))
    print(f"  cs_rank(close / volume) -> {expr4}")

    # 测试 5: 复杂表达式
    print("\n测试 5: 复杂表达式")
    ma20 = FunctionNode("ts_mean", close, 20)
    ma60 = FunctionNode("ts_mean", close, 60)
    expr5 = translator.translate(FunctionNode("cs_rank", ma20 - ma60))
    print(f"  cs_rank(ts_mean(close, 20) - ts_mean(close, 60)) -> {expr5}")

    print("\n✅ 所有测试通过!")
