"""函数装饰器 - 用于标记因子函数类型并自动注入 over 处理。

提供三种装饰器：
- @time_series: 时序因子，自动添加 .over("ts_code")
- @cross_section: 截面因子，自动添加 .over("trade_date")
- @element_wise: 元素级运算，不添加 over
"""

from functools import wraps
from typing import Callable


def time_series(func: Callable) -> Callable:
    """标记为时序因子，自动添加 .over('ts_code')。

    用于 ts_* 函数，如 ts_mean, ts_std, ts_corr 等。
    每个时序计算都按股票代码分组，防止跨股票串数据。
    """

    @wraps(func)
    def wrapper(self, node):
        expr = func(self, node)
        return expr.over("ts_code")

    wrapper._factor_type = "ts"
    return wrapper


def cross_section(func: Callable) -> Callable:
    """标记为截面因子，自动添加 .over('trade_date')。

    用于 cs_* 函数，如 cs_rank, cs_zscore 等。
    每个截面计算都按交易日分组，在同一天的所有股票间计算。
    """

    @wraps(func)
    def wrapper(self, node):
        expr = func(self, node)
        return expr.over("trade_date")

    wrapper._factor_type = "cs"
    return wrapper


def element_wise(func: Callable) -> Callable:
    """标记为元素级运算，不添加 over。

    用于数学函数，如 log, exp, sqrt, cos, sin 等。
    这些函数对每个元素独立计算，不需要分组。
    """

    @wraps(func)
    def wrapper(self, node):
        return func(self, node)

    wrapper._factor_type = "element"
    return wrapper