"""用户可配置的股票池定义文件。

参考 `src/experiment/common.py` 中的 `stock_pool_filter` 模式，
支持通过编写 Polars filter 函数来自定义股票池。
"""

from typing import Callable, List

import polars as pl

from src.factorminer.evaluation.stock_pool_registry import StockPoolDefinition, StockPoolRegistry


# =============================================================================
# 股票池大小配置
# =============================================================================
SMALL_CAP_N = 1000  # 小微盘股票池默认选取的股票数量


# =============================================================================
# 股票池筛选函数（与 experiment/common.py 风格一致）
# =============================================================================
def all_market_filter(df: pl.DataFrame) -> pl.Series:
    """全市场股票池。"""
    return pl.Series([True] * len(df))


def growth_board_filter(df: pl.DataFrame) -> pl.Series:
    """创业板股票池（代码以 300 开头）。"""
    return df["ts_code"].str.starts_with("300")


def star_board_filter(df: pl.DataFrame) -> pl.Series:
    """科创板股票池（代码以 688 开头）。"""
    return df["ts_code"].str.starts_with("688")


def bse_board_filter(df: pl.DataFrame) -> pl.Series:
    """北交所股票池（代码以 8 或 4 开头）。"""
    return df["ts_code"].str.starts_with("8") | df["ts_code"].str.starts_with("4")


def main_board_filter(df: pl.DataFrame) -> pl.Series:
    """主板股票池（排除创业板、科创板、北交所）。"""
    return (
        ~df["ts_code"].str.starts_with("300")
        & ~df["ts_code"].str.starts_with("688")
        & ~df["ts_code"].str.starts_with("8")
        & ~df["ts_code"].str.starts_with("4")
    )


def small_cap_filter(df: pl.DataFrame, n_stocks: int = SMALL_CAP_N) -> pl.Series:
    """小微盘股票池（市值最小的 n_stocks 只股票）。

    Args:
        df: 包含 ts_code 和 circ_mv 列的数据框。
        n_stocks: 选取的股票数量，默认 SMALL_CAP_N。

    Returns:
        布尔 Series，表示是否入选小微盘股票池。
    """
    if "circ_mv" not in df.columns:
        # 若缺失 circ_mv 数据则全部排除（安全降级）
        return pl.Series([False] * len(df))
    n = min(n_stocks, len(df))
    small_codes = df.sort("circ_mv").head(n)["ts_code"]
    return df["ts_code"].is_in(small_codes.implode())


def main_board_small_cap_filter(
    df: pl.DataFrame, n_stocks: int = SMALL_CAP_N
) -> pl.Series:
    """主板小微盘股票池（主板中市值最小的 n_stocks 只股票）。"""
    main_board = main_board_filter(df)
    main_df = df.filter(main_board)
    if "circ_mv" not in df.columns or len(main_df) == 0:
        return pl.Series([False] * len(df))
    n = min(n_stocks, len(main_df))
    small_codes = main_df.sort("circ_mv").head(n)["ts_code"]
    return df["ts_code"].is_in(small_codes.implode())


def common_small_cap_filter(df: pl.DataFrame, n_stocks: int = SMALL_CAP_N) -> pl.Series:
    """与 experiment/common.py 完全一致的小微盘筛选函数。

    筛选条件：
    1. 排除创业板（代码以 300 开头）
    2. 排除科创板（代码以 688 开头）
    3. 排除北交所（代码以 8、9 或 4 开头）
    4. 选取当日流通市值最小的 n_stocks 只股票

    Args:
        df: 数据框，必须包含 ts_code 和 circ_mv 列
        n_stocks: 选取的股票数量，默认 SMALL_CAP_N

    Returns:
        布尔 Series，表示哪些股票被选中
    """
    code_filter = (
        ~df["ts_code"].str.starts_with("30")  # 排除创业板
        & ~df["ts_code"].str.starts_with("68")  # 排除科创板
        & ~df["ts_code"].str.starts_with("8")  # 排除北交所
        & ~df["ts_code"].str.starts_with("9")  # 排除北交所
        & ~df["ts_code"].str.starts_with("4")  # 排除北交所
    )
    valid_df = df.filter(code_filter)
    n = min(n_stocks, len(valid_df))
    small_codes = valid_df.sort("circ_mv").head(n)["ts_code"]
    return df["ts_code"].is_in(small_codes.implode())


# =============================================================================
# 默认股票池定义列表
# =============================================================================
def get_default_stock_pool_definitions() -> List[StockPoolDefinition]:
    """返回默认的股票池定义列表。"""
    return [
        StockPoolDefinition(
            name="all",
            filter_func=all_market_filter,
            required_columns=[],
            description="全市场",
        ),
        StockPoolDefinition(
            name="growth_board",
            filter_func=growth_board_filter,
            required_columns=[],
            description="创业板",
        ),
        StockPoolDefinition(
            name="star_board",
            filter_func=star_board_filter,
            required_columns=[],
            description="科创板",
        ),
        StockPoolDefinition(
            name="bse_board",
            filter_func=bse_board_filter,
            required_columns=[],
            description="北交所",
        ),
        StockPoolDefinition(
            name="main_board",
            filter_func=main_board_filter,
            required_columns=[],
            description="主板",
        ),
        StockPoolDefinition(
            name="small_cap",
            filter_func=small_cap_filter,
            required_columns=["circ_mv"],
            description="小微盘",
        ),
        StockPoolDefinition(
            name="main_board_small_cap",
            filter_func=main_board_small_cap_filter,
            required_columns=["circ_mv"],
            description="主板小微盘",
        ),
    ]


def get_default_stock_pools(
    pre_filters: List[Callable[[pl.DataFrame], pl.Series | pl.Expr]] | None = None,
) -> StockPoolRegistry:
    """返回默认的股票池注册表。"""
    return StockPoolRegistry.from_definitions(
        get_default_stock_pool_definitions(),
        pre_filters=pre_filters,
    )
