"""股票池注册表：支持配置化股票池筛选与掩码生成。"""

from dataclasses import dataclass, field
from typing import Callable, Dict, List, Optional

import numpy as np
import polars as pl


@dataclass
class StockPoolDefinition:
    """股票池定义，参考 experiment/common.py 中的 filter 模式。"""

    name: str
    filter_func: Callable[[pl.DataFrame], pl.Series | pl.Expr]
    required_columns: List[str] = field(default_factory=list)
    description: str = ""


class StockPoolRegistry:
    """管理多个股票池的定义，并为 (M,) 资产列表生成布尔掩码。"""

    def __init__(
        self,
        definitions: Optional[List[StockPoolDefinition]] = None,
        pre_filters: Optional[
            List[Callable[[pl.DataFrame], pl.Series | pl.Expr]]
        ] = None,
    ) -> None:
        self.pre_filters: List[Callable[[pl.DataFrame], pl.Series | pl.Expr]] = (
            pre_filters or []
        )
        self.pools: Dict[str, dict] = {}
        self.masks: Dict[str, np.ndarray] = {}
        self._resolved = False
        if definitions:
            for d in definitions:
                self.add_pool(d.name, d.filter_func, d.required_columns)

    @classmethod
    def from_definitions(
        cls,
        definitions: List[StockPoolDefinition],
        pre_filters: Optional[
            List[Callable[[pl.DataFrame], pl.Series | pl.Expr]]
        ] = None,
    ) -> "StockPoolRegistry":
        """通过定义列表快速创建注册表。"""
        return cls(definitions=definitions, pre_filters=pre_filters)

    def add_pool(
        self,
        name: str,
        filter_func: Callable[[pl.DataFrame], pl.Series | pl.Expr],
        required_columns: Optional[List[str]] = None,
    ) -> None:
        """注册一个股票池。

        Args:
            name: 股票池名称，如 "small_cap"。
            filter_func: 接收包含 ts_code 和 required_columns 的 DataFrame，
                返回布尔 Series 或 Expr。
            required_columns: filter_func 额外需要的列名列表。
        """
        self.pools[name] = {
            "filter_func": filter_func,
            "required_columns": required_columns or [],
        }
        self._resolved = False

    @staticmethod
    def _eval_filter_result(
        df: pl.DataFrame,
        result: pl.Series | pl.Expr,
        context: str,
    ) -> np.ndarray:
        """将 filter_func 的返回值统一转换为 numpy bool 数组。"""
        if isinstance(result, pl.Series):
            mask_series = result
        elif isinstance(result, pl.Expr):
            mask_series = df.select(result.alias("_mask")).to_series()
        else:
            raise TypeError(
                f"{context} 的 filter_func 必须返回 pl.Series 或 pl.Expr，"
                f"实际返回 {type(result)}"
            )
        return mask_series.to_numpy().astype(bool)

    def build_masks(
        self,
        asset_codes: List[str],
        metadata_df: Optional[pl.DataFrame] = None,
    ) -> Dict[str, np.ndarray]:
        """为所有注册的股票池生成布尔掩码。

        Args:
            asset_codes: 按矩阵顺序排列的股票代码列表 (M,)。
            metadata_df: 包含 extra columns 的 metadata，可选。

        Returns:
            {pool_name: bool_array_of_shape_(M,)} 字典。
        """
        base_df = pl.DataFrame({"ts_code": asset_codes})
        df = base_df
        if metadata_df is not None:
            df = base_df.join(metadata_df, on="ts_code", how="left")

        base_mask = np.ones(len(asset_codes), dtype=bool)
        for pre_filter in self.pre_filters:
            result = pre_filter(df)
            base_mask &= self._eval_filter_result(df, result, "pre_filter")

        self.masks = {}
        for name, cfg in self.pools.items():
            result = cfg["filter_func"](df)
            pool_mask = self._eval_filter_result(df, result, f"股票池 '{name}'")
            self.masks[name] = base_mask & pool_mask

        self._resolved = True
        return self.masks

    def filter_signals(self, signals: np.ndarray, pool_name: str) -> np.ndarray:
        """使用已构建的掩码过滤信号矩阵。

        Args:
            signals: (M, T) 信号矩阵。
            pool_name: 股票池名称。

        Returns:
            (M_pool, T) 的子矩阵。
        """
        if not self._resolved:
            raise RuntimeError("请先调用 build_masks 构建掩码")
        mask = self.masks.get(pool_name)
        if mask is None:
            raise KeyError(f"未知的股票池: {pool_name}")
        return signals[mask, :]

    def get_pool_names(self) -> List[str]:
        """返回已注册的股票池名称列表。"""
        return list(self.pools.keys())

    def get_required_columns(self) -> List[str]:
        """返回所有 pre_filter 和股票池所需的列名并集。"""
        cols: set[str] = set()
        for pre_filter in self.pre_filters:
            cols.update(getattr(pre_filter, "required_columns", []))
        for cfg in self.pools.values():
            cols.update(cfg["required_columns"])
        return sorted(cols)
