ProStock/src/factors/engine.py

"""执行引擎 - Phase 4 因子执行引擎

本模块负责因子计算的核心逻辑：
- FactorEngine: 因子执行引擎，根据因子类型采用不同的计算和防泄露策略

防泄露策略：
1. CrossSectionalFactor：防止日期泄露，每天传入 [T-lookback+1, T] 数据
2. TimeSeriesFactor：防止股票泄露，每只股票传入完整序列
"""

from typing import List, Optional

import polars as pl

from src.factors.data_loader import DataLoader
from src.factors.data_spec import FactorContext, FactorData
from src.factors.base import BaseFactor, CrossSectionalFactor, TimeSeriesFactor


class FactorEngine:
    """因子执行引擎 - 根据因子类型采用不同的计算和防泄露策略

    核心职责：
    1. CrossSectionalFactor：防止日期泄露，每天传入 [T-lookback+1, T] 数据
    2. TimeSeriesFactor：防止股票泄露，每只股票传入完整序列

    示例：
        >>> loader = DataLoader(data_dir="data")
        >>> engine = FactorEngine(loader)
        >>> result = engine.compute(factor, start_date="20240101", end_date="20240131")
    """

    def __init__(self, data_loader: DataLoader):
        """初始化引擎

        Args:
            data_loader: 数据加载器实例
        """
        self.data_loader = data_loader

    def compute(self, factor: BaseFactor, **kwargs) -> pl.DataFrame:
        """统一的计算入口

        根据 factor_type 分发到具体方法：
        - "cross_sectional" -> _compute_cross_sectional()
        - "time_series" -> _compute_time_series()

        Args:
            factor: 要计算的因子
            **kwargs: 额外参数，根据因子类型不同：
                - 截面因子: start_date, end_date
                - 时序因子: stock_codes, start_date, end_date

        Returns:
            DataFrame[trade_date, ts_code, factor_name]

        Raises:
            ValueError: 无效的 factor_type 或缺少必需参数
        """
        if factor.factor_type == "cross_sectional":
            if "start_date" not in kwargs or "end_date" not in kwargs:
                raise ValueError(
                    "cross_sectional factor requires 'start_date' and 'end_date' parameters"
                )
            return self._compute_cross_sectional(
                factor, kwargs["start_date"], kwargs["end_date"]
            )
        elif factor.factor_type == "time_series":
            missing = []
            if "stock_codes" not in kwargs:
                missing.append("stock_codes")
            if "start_date" not in kwargs:
                missing.append("start_date")
            if "end_date" not in kwargs:
                missing.append("end_date")
            if missing:
                raise ValueError(
                    f"time_series factor requires parameters: {', '.join(missing)}"
                )
            return self._compute_time_series(
                factor,
                kwargs["stock_codes"],
                kwargs["start_date"],
                kwargs["end_date"],
            )
        else:
            raise ValueError(f"Unknown factor type: {factor.factor_type}")

    def _compute_cross_sectional(
        self,
        factor: CrossSectionalFactor,
        start_date: str,
        end_date: str,
    ) -> pl.DataFrame:
        """执行日期截面计算

        防泄露策略：
        - 防止日期泄露：每天只传入 [T-lookback+1, T] 的数据（不含未来）
        - 允许股票间比较：传入当天所有股票的数据

        计算流程：
        1. 计算 max_lookback，确定数据起始日期
        2. 一次性加载 [start-max_lookback+1, end] 的所有数据
        3. 对每个日期 T in [start_date, end_date]：
           a. 裁剪数据到 [T-lookback+1, T]
           b. 创建 FactorData（current_date=T）
           c. 调用 factor.compute()
           d. 收集结果
        4. 合并所有日期的结果

        返回 DataFrame 格式：
        ┌────────────┬──────────┬──────────────┐
        │ trade_date │ ts_code  │ factor_name  │
        ├────────────┼──────────┼──────────────┤
        │ 20240101   │ 000001.SZ│ 0.5          │
        │ 20240101   │ 000002.SZ│ 0.3          │
        └────────────┴──────────┴──────────────┘

        Args:
            factor: 截面因子
            start_date: 开始日期 YYYYMMDD
            end_date: 结束日期 YYYYMMDD

        Returns:
            包含因子值的 DataFrame
        """
        # 计算最大 lookback
        max_lookback = max(spec.lookback_days for spec in factor.data_specs)

        # 确定数据起始日期（向前扩展 lookback）
        data_start = self._get_trading_date_offset(start_date, -max_lookback + 1)

        # 一次性加载所有数据
        raw_data = self.data_loader.load(
            factor.data_specs, date_range=(data_start, end_date)
        )

        results = []

        # 获取日期范围
        date_range = self._get_date_range(start_date, end_date, raw_data)

        # 按日期遍历：每天计算一次
        for current_date in date_range:
            # 裁剪数据：只保留 current_date 及之前的数据（防止日期泄露）
            # 但保留所有股票的数据（允许股票间比较）
            day_data = raw_data.filter(pl.col("trade_date") <= current_date)

            # 如果 lookback > 0，进一步裁剪到 lookback 窗口
            if max_lookback > 0:
                lookback_start = self._get_trading_date_offset(
                    current_date, -max_lookback + 1
                )
                day_data = day_data.filter(pl.col("trade_date") >= lookback_start)

            # 如果没有数据，跳过
            if len(day_data) == 0:
                continue

            # 创建 FactorData（包含当天及历史数据，无未来数据）
            context = FactorContext(
                current_date=current_date,
                trade_dates=date_range,
            )
            factor_data = FactorData(day_data, context)

            # 计算因子值
            factor_values = factor.compute(factor_data)

            # 获取当前日期的股票列表
            today_stocks = day_data.filter(pl.col("trade_date") == current_date)[
                "ts_code"
            ]

            # 确保 factor_values 长度与股票列表一致
            if len(factor_values) != len(today_stocks):
                # 如果长度不一致，可能是 factor.compute 返回了错误的长度
                # 尝试从 factor_data 重新提取
                cs_data = factor_data.get_cross_section()
                if len(cs_data) > 0:
                    today_stocks = cs_data["ts_code"]
                    # 如果 factor_values 仍然不匹配，用 null 填充
                    if len(factor_values) != len(today_stocks):
                        factor_values = pl.Series([None] * len(today_stocks))

            # 收集结果
            if len(today_stocks) > 0:
                results.append(
                    pl.DataFrame(
                        {
                            "trade_date": [current_date] * len(today_stocks),
                            "ts_code": today_stocks,
                            factor.name: factor_values,
                        }
                    )
                )

        # 合并所有日期的结果
        if results:
            return pl.concat(results)
        else:
            # 返回空 DataFrame
            return pl.DataFrame(
                {
                    "trade_date": [],
                    "ts_code": [],
                    factor.name: [],
                }
            )

    def _compute_time_series(
        self,
        factor: TimeSeriesFactor,
        stock_codes: List[str],
        start_date: str,
        end_date: str,
    ) -> pl.DataFrame:
        """执行时间序列计算

        防泄露策略：
        - 防止股票泄露：每只股票单独计算，传入该股票的完整序列
        - 允许访问历史数据：时序计算需要历史数据，这是正常的

        计算流程：
        1. 计算 max_lookback，确定数据起始日期
        2. 一次性加载 [start-max_lookback+1, end] 的所有数据
        3. 对每只股票 S in stock_codes：
           a. 过滤出 S 的数据（防止股票泄露）
           b. 创建 FactorData（current_stock=S）
           c. 调用 factor.compute()（向量化计算整个序列）
           d. 收集结果
        4. 合并所有股票的结果

        性能优势：
        - 使用 Polars 的 rolling_mean 等向量化操作
        - 每只股票只计算一次，无重复计算

        返回 DataFrame 格式：
        ┌────────────┬──────────┬──────────────┐
        │ trade_date │ ts_code  │ factor_name  │
        ├────────────┼──────────┼──────────────┤
        │ 20240101   │ 000001.SZ│ 10.5         │
        │ 20240102   │ 000001.SZ│ 10.6         │
        └────────────┴──────────┴──────────────┘

        Args:
            factor: 时序因子
            stock_codes: 股票代码列表
            start_date: 开始日期 YYYYMMDD
            end_date: 结束日期 YYYYMMDD

        Returns:
            包含因子值的 DataFrame
        """
        # 计算最大 lookback
        max_lookback = max(spec.lookback_days for spec in factor.data_specs)

        # 确定数据起始日期（向前扩展 lookback）
        data_start = self._get_trading_date_offset(start_date, -max_lookback + 1)

        # 加载所有数据
        all_data = self.data_loader.load(
            factor.data_specs, date_range=(data_start, end_date)
        )

        results = []

        # 获取所有交易日
        all_dates = all_data["trade_date"].unique().sort() if len(all_data) > 0 else []

        # 按股票遍历：每只股票一次性计算
        for stock_code in stock_codes:
            # 过滤出该股票的数据（防止股票泄露）
            stock_data = all_data.filter(pl.col("ts_code") == stock_code)

            if len(stock_data) == 0:
                continue

            # 创建 FactorData（该股票的完整序列）
            context = FactorContext(
                current_stock=stock_code,
                trade_dates=list(all_dates),
            )
            factor_data = FactorData(stock_data, context)

            # 一次性计算整个时间序列（向量化，高效）
            factor_values = factor.compute(factor_data)

            # 获取该股票的日期列表
            stock_dates = stock_data["trade_date"]

            # 确保 factor_values 长度与日期列表一致
            if len(factor_values) != len(stock_dates):
                # 如果长度不一致，用 null 填充
                factor_values = pl.Series([None] * len(stock_dates))

            # 收集结果
            results.append(
                pl.DataFrame(
                    {
                        "trade_date": stock_dates,
                        "ts_code": [stock_code] * len(stock_dates),
                        factor.name: factor_values,
                    }
                )
            )

        # 合并所有股票的结果
        if results:
            return pl.concat(results)
        else:
            # 返回空 DataFrame
            return pl.DataFrame(
                {
                    "trade_date": [],
                    "ts_code": [],
                    factor.name: [],
                }
            )

    def _get_trading_date_offset(self, date: str, offset: int) -> str:
        """获取相对于给定日期的交易日偏移

        简单实现：假设每天都有交易，直接计算日期偏移
        实际项目中可能需要使用交易日历

        Args:
            date: 基准日期 YYYYMMDD
            offset: 偏移天数（正数向后，负数向前）

        Returns:
            偏移后的日期 YYYYMMDD
        """
        from datetime import datetime, timedelta

        dt = datetime.strptime(date, "%Y%m%d")
        new_dt = dt + timedelta(days=offset)
        return new_dt.strftime("%Y%m%d")

    def _get_date_range(
        self, start_date: str, end_date: str, data: pl.DataFrame
    ) -> List[str]:
        """获取日期范围内的所有交易日

        Args:
            start_date: 开始日期 YYYYMMDD
            end_date: 结束日期 YYYYMMDD
            data: 包含 trade_date 列的 DataFrame

        Returns:
            日期列表
        """
        if len(data) == 0:
            return []

        # 从数据中获取实际存在的日期
        dates = (
            data.filter(
                (pl.col("trade_date") >= start_date)
                & (pl.col("trade_date") <= end_date)
            )["trade_date"]
            .unique()
            .sort()
            .to_list()
        )

        return dates