"""股票池管理器

支持使用自定义函数和因子表达式进行每日股票池筛选。
临时计算的因子仅在筛选阶段使用，绝不泄露到训练数据。
"""

from typing import TYPE_CHECKING, Callable, Dict, List, Optional, Set, Tuple

import polars as pl

if TYPE_CHECKING:
    from src.factors.engine.data_router import DataRouter


class StockPoolManager:
    """股票池管理器 - 支持自定义筛选函数和因子

    核心特性：
    1. 支持传入自定义筛选函数
    2. 支持使用因子表达式进行筛选
    3. 使用 FactorEngine 计算所需因子
    4. 只删除本次新生成的临时因子，保留输入中已存在的所有列

    数据流：
        输入数据 (含原始列，可能包含一些因子)
            ↓
        [准备数据]
            ├─ 获取缺失的基础列 (from data_router)
            └─ 计算缺失的因子 (使用 FactorEngine，标记为"本次生成")
            ↓
        [每日筛选]
            ├─ group_by("trade_date").apply(filter_func)
            └─ 只保留 ts_code + trade_date (筛选结果标识)
            ↓
        [返回结果]
            └─ semi join 原始数据，保留所有原始列
    """

    def __init__(
        self,
        filter_func: Callable[[pl.DataFrame], pl.Series],
        required_columns: Optional[List[str]] = None,
        required_factors: Optional[Dict[str, str]] = None,
        data_router: Optional["DataRouter"] = None,
        code_col: str = "ts_code",
        date_col: str = "trade_date",
    ):
        """初始化股票池管理器

        Args:
            filter_func: 筛选函数，接收 DataFrame 返回布尔 Series
            required_columns: 除输入数据外还需获取的基础列
            required_factors: 筛选所需的因子表达式 {因子名: DSL表达式}
            data_router: 数据路由器，用于获取缺失列
            code_col: 股票代码列名
            date_col: 日期列名
        """
        self.filter_func = filter_func
        self.required_columns = required_columns or []
        self.required_factors = required_factors or {}
        self.data_router = data_router
        self.code_col = code_col
        self.date_col = date_col

    def filter_and_select_daily(self, data: pl.DataFrame) -> pl.DataFrame:
        """每日独立筛选股票池

        流程：
        1. 记录输入数据的原始列
        2. 收集筛选所需的完整数据（基础列 + 计算因子）
        3. 按日期分组应用筛选函数
        4. 只返回 ts_code 和 trade_date（筛选结果标识）
        5. 用标识列从原始数据筛选（保留所有原始列）

        关键：返回的数据包含输入数据的所有原始列，只移除本次新生成的临时因子

        Args:
            data: 因子计算后的全市场数据，必须包含 trade_date 和 ts_code 列

        Returns:
            筛选后的数据，列与输入数据完全一致（临时因子已移除）
        """
        # 1. 记录原始列，用于最后验证
        original_columns = list(data.columns)

        # 2. 准备完整数据（用于筛选判断）
        # 返回的 enriched 包含临时因子，但不修改原始 data
        enriched = self._prepare_data(data)

        # 3. 每日筛选，只保留标识列
        # 使用 group_by + map_groups 替代 apply（Polars 0.20+）
        selected_ids = enriched.group_by(self.date_col).map_groups(
            lambda df: df.filter(self.filter_func(df)).select(
                [self.code_col, self.date_col]
            )
        )

        # 4. 用 semi join 从原始数据筛选，自动只保留原始列
        # semi join: 保留左侧（data）的所有列，只保留匹配的行
        result = data.join(
            selected_ids,
            on=[self.code_col, self.date_col],
            how="semi",
        )

        # 5. 验证：确保结果列与原始列完全一致
        if list(result.columns) != original_columns:
            raise RuntimeError(
                f"列发生变化！\n原始: {original_columns}\n结果: {list(result.columns)}"
            )

        return result

    def _prepare_data(self, data: pl.DataFrame) -> pl.DataFrame:
        """准备筛选所需的完整数据

        步骤：
        1. 获取缺失的基础列
        2. 计算缺失的因子（输入中已存在的因子不再计算）

        Args:
            data: 输入数据

        Returns:
            包含所有所需列和因子的数据（含临时因子）
        """
        result = data

        # 1. 获取缺失的基础列
        if self.required_columns and self.data_router is not None:
            result = self._fetch_required_columns(result)

        # 2. 计算因子（只计算输入中不存在的）
        if self.required_factors:
            result = self._compute_factors(result)

        return result

    def _fetch_required_columns(self, data: pl.DataFrame) -> pl.DataFrame:
        """从 data_router 获取缺失的基础列

        Args:
            data: 当前数据

        Returns:
            补充了缺失列的数据
        """
        missing_cols = set(self.required_columns) - set(data.columns)
        if not missing_cols:
            return data

        if self.data_router is None:
            raise ValueError(f"需要获取列 {missing_cols}，但未提供 data_router")

        # 获取日期范围
        dates = data.select(self.date_col).unique().to_series().to_list()
        if not dates:
            return data

        start_date = min(dates)
        end_date = max(dates)

        # 获取所有股票代码
        codes = data.select(self.code_col).unique().to_series().to_list()

        try:
            from src.factors.engine.data_spec import DataSpec

            # 构建 DataSpec 列表
            data_specs = [
                DataSpec("daily", list(missing_cols))  # 假设从 daily 表获取
            ]

            # 从 data_router 获取数据
            extra_data = self.data_router.fetch_data(
                data_specs=data_specs,
                start_date=start_date,
                end_date=end_date,
                stock_codes=codes,
            )

            # 合并到结果
            result = data.join(
                extra_data,
                on=[self.code_col, self.date_col],
                how="left",
            )

            return result

        except Exception as e:
            print(f"[警告] 获取缺失列失败: {e}")
            # 如果获取失败，继续使用现有数据（筛选可能不完全）
            return data

    def _compute_factors(self, data: pl.DataFrame) -> pl.DataFrame:
        """使用 FactorEngine 计算筛选所需的因子

        只计算输入数据中不存在的因子，已存在的因子直接使用。

        Args:
            data: 当前数据

        Returns:
            补充了缺失因子的数据（含临时因子）
        """
        existing_cols = set(data.columns)

        # 确定需要计算的因子（输入中不存在的）
        factors_to_compute = {
            name: expr
            for name, expr in self.required_factors.items()
            if name not in existing_cols
        }

        if not factors_to_compute:
            # 所有因子都已存在，无需计算
            return data

        try:
            from src.factors import FactorEngine

            # 获取日期范围
            dates = data.select(self.date_col).unique().to_series().to_list()
            if not dates:
                return data

            start_date = min(dates)
            end_date = max(dates)

            # 创建 FactorEngine 并注册因子
            engine = FactorEngine()
            for name, expr in factors_to_compute.items():
                engine.add_factor(name, expr)

            # 计算因子
            factor_data = engine.compute(
                factor_names=list(factors_to_compute.keys()),
                start_date=start_date,
                end_date=end_date,
            )

            # 合并到数据（左连接，保留所有原始行）
            result = data.join(
                factor_data,
                on=[self.code_col, self.date_col],
                how="left",
            )

            return result

        except Exception as e:
            print(f"[警告] 计算因子失败: {e}")
            # 如果计算失败，继续使用现有数据
            return data
