"""测试 StockPoolManager

验证新的自定义函数和因子筛选功能。
重点测试：临时因子隔离（只删除新生成的因子，保留原本存在的）。
"""

from unittest.mock import Mock, patch

import polars as pl
import pytest

from src.training.core.stock_pool_manager import StockPoolManager


class TestStockPoolManagerBasic:
    """StockPoolManager 基础测试类"""

    def test_basic_filter_with_columns(self):
        """测试使用基础列进行筛选"""

        def filter_func(df: pl.DataFrame) -> pl.Series:
            return df["total_mv"] > 50

        # 创建模拟 data_router
        mock_router = Mock()

        manager = StockPoolManager(
            filter_func=filter_func,
            required_columns=["total_mv"],
            data_router=mock_router,
        )

        # 创建测试数据
        data = pl.DataFrame(
            {
                "ts_code": ["000001.SZ", "000002.SZ", "000003.SZ", "000004.SZ"],
                "trade_date": ["20240101"] * 4,
                "close": [10.0, 20.0, 30.0, 40.0],
                "total_mv": [100.0, 30.0, 80.0, 20.0],
            }
        )

        # 执行筛选（无需 mock，因为 total_mv 已在数据中）
        result = manager.filter_and_select_daily(data)

        # 验证返回数据列与输入一致
        assert result.columns == data.columns
        # 验证筛选生效（保留市值 > 50 的股票）
        assert len(result) == 2
        assert "000001.SZ" in result["ts_code"].to_list()
        assert "000003.SZ" in result["ts_code"].to_list()

    def test_filter_without_required_columns(self):
        """测试不使用额外列，仅使用输入数据中已有的列"""

        def filter_func(df: pl.DataFrame) -> pl.Series:
            return df["close"] > 25

        manager = StockPoolManager(filter_func=filter_func)

        data = pl.DataFrame(
            {
                "ts_code": ["000001.SZ", "000002.SZ", "000003.SZ"],
                "trade_date": ["20240101"] * 3,
                "close": [10.0, 30.0, 20.0],
            }
        )

        result = manager.filter_and_select_daily(data)

        # 验证只保留 close > 25 的股票
        assert len(result) == 1
        assert result["ts_code"][0] == "000002.SZ"
        assert result.columns == data.columns

    def test_empty_result(self):
        """测试筛选结果为空的情况"""

        def filter_func(df: pl.DataFrame) -> pl.Series:
            return df["close"] > 9999  # 不可能满足的条件

        manager = StockPoolManager(filter_func=filter_func)

        data = pl.DataFrame(
            {
                "ts_code": ["000001.SZ", "000002.SZ"],
                "trade_date": ["20240101"] * 2,
                "close": [10.0, 20.0],
            }
        )

        result = manager.filter_and_select_daily(data)

        assert len(result) == 0
        assert result.columns == data.columns  # 即使为空，列结构保持一致


class TestStockPoolManagerDailyIndependence:
    """每日独立筛选测试类"""

    def test_daily_independence(self):
        """测试每日独立进行筛选"""

        def filter_func(df: pl.DataFrame) -> pl.Series:
            # 每日选收盘价前 50%
            median = df["close"].median()
            return df["close"] >= median

        manager = StockPoolManager(filter_func=filter_func)

        # 创建多日期数据
        data = pl.DataFrame(
            {
                "ts_code": [
                    "000001.SZ",
                    "000002.SZ",
                    "000003.SZ",
                    "000004.SZ",
                    # 日期 2
                    "000001.SZ",
                    "000002.SZ",
                    "000003.SZ",
                    "000004.SZ",
                ],
                "trade_date": [
                    "20240101",
                    "20240101",
                    "20240101",
                    "20240101",
                    "20240102",
                    "20240102",
                    "20240102",
                    "20240102",
                ],
                "close": [
                    10.0,
                    20.0,
                    30.0,
                    40.0,  # 日期1：选 30, 40
                    5.0,
                    15.0,
                    25.0,
                    35.0,  # 日期2：选 25, 35
                ],
            }
        )

        result = manager.filter_and_select_daily(data)

        # 验证每个日期独立筛选
        day1 = result.filter(pl.col("trade_date") == "20240101")
        day2 = result.filter(pl.col("trade_date") == "20240102")

        # 日期1：收盘价 >= 25（中位数）- 30 和 40
        assert len(day1) == 2
        assert set(day1["ts_code"].to_list()) == {"000003.SZ", "000004.SZ"}

        # 日期2：收盘价 >= 20（中位数）- 25 和 35
        assert len(day2) == 2
        assert set(day2["ts_code"].to_list()) == {"000003.SZ", "000004.SZ"}

    def test_uneven_daily_distribution(self):
        """测试每日股票数量不均的情况"""

        def filter_func(df: pl.DataFrame) -> pl.Series:
            return df["close"] > 15

        manager = StockPoolManager(filter_func=filter_func)

        data = pl.DataFrame(
            {
                "ts_code": [
                    "000001.SZ",
                    "000002.SZ",  # 日期1：2只股票
                    "000001.SZ",
                    "000002.SZ",
                    "000003.SZ",
                    "000004.SZ",  # 日期2：4只股票
                ],
                "trade_date": [
                    "20240101",
                    "20240101",
                    "20240102",
                    "20240102",
                    "20240102",
                    "20240102",
                ],
                "close": [10.0, 20.0, 5.0, 15.0, 25.0, 35.0],
            }
        )

        result = manager.filter_and_select_daily(data)

        # 日期1：只有 000002.SZ (20 > 15)
        day1 = result.filter(pl.col("trade_date") == "20240101")
        assert len(day1) == 1

        # 日期2：000003.SZ (25) 和 000004.SZ (35)
        day2 = result.filter(pl.col("trade_date") == "20240102")
        assert len(day2) == 2


class TestStockPoolManagerFactorIsolation:
    """因子隔离测试类 - 核心测试"""

    @patch.object(StockPoolManager, "_compute_factors")
    def test_filter_with_factors(self, mock_compute):
        """测试使用因子表达式进行筛选，验证临时因子被删除"""

        # 设置 mock 返回值（包含计算后的因子）
        mock_compute.return_value = pl.DataFrame(
            {
                "ts_code": ["000001.SZ", "000002.SZ", "000003.SZ"],
                "trade_date": ["20240101"] * 3,
                "close": [11.0, 9.5, 10.8],
                "momentum_20": [0.1, -0.05, 0.08],  # 只有第一个 > 0.05
            }
        )

        # 创建 Manager
        def filter_func(df: pl.DataFrame) -> pl.Series:
            return df["momentum_20"] > 0.05

        manager = StockPoolManager(
            filter_func=filter_func,
            required_factors={"momentum_20": "(close / ts_delay(close, 20)) - 1"},
        )

        # 输入数据不含 momentum_20
        data = pl.DataFrame(
            {
                "ts_code": ["000001.SZ", "000002.SZ", "000003.SZ"],
                "trade_date": ["20240101"] * 3,
                "close": [11.0, 9.5, 10.8],
            }
        )

        result = manager.filter_and_select_daily(data)

        # 验证返回数据列与输入一致（momentum_20 被删除）
        assert result.columns == data.columns
        assert "momentum_20" not in result.columns

        # 验证筛选生效
        # momentum_20 > 0.05: 000001.SZ (0.1), 000003.SZ (0.08)
        assert len(result) == 2
        assert "000001.SZ" in result["ts_code"].to_list()
        assert "000003.SZ" in result["ts_code"].to_list()

        # 验证 _compute_factors 被调用
        mock_compute.assert_called_once()

    @patch.object(StockPoolManager, "_compute_factors")
    def test_preserve_existing_factors(self, mock_compute):
        """测试输入中已存在的因子不会被删除（核心测试）"""

        # 设置 mock 返回值（包含 roe，但 momentum_20 已在输入中）
        mock_compute.return_value = pl.DataFrame(
            {
                "ts_code": ["000001.SZ", "000002.SZ", "000003.SZ"],
                "trade_date": ["20240101"] * 3,
                "close": [11.0, 9.5, 10.8],
                "momentum_20": [0.1, -0.05, 0.08],  # 原本就存在
                "roe": [0.12, 0.08, 0.15],  # 本次生成，第二个 < 0.1
            }
        )

        def filter_func(df: pl.DataFrame) -> pl.Series:
            return (df["momentum_20"] > 0.05) & (df["roe"] > 0.1)

        manager = StockPoolManager(
            filter_func=filter_func,
            # 只声明 roe 为本次生成，momentum_20 已在输入中
            required_factors={"roe": "n_income / equity"},
        )

        # 输入数据已包含 momentum_20
        data = pl.DataFrame(
            {
                "ts_code": ["000001.SZ", "000002.SZ", "000003.SZ"],
                "trade_date": ["20240101"] * 3,
                "close": [11.0, 9.5, 10.8],
                "momentum_20": [0.1, -0.05, 0.08],  # 原本就存在的因子
            }
        )

        result = manager.filter_and_select_daily(data)

        # 关键断言：
        # 1. momentum_20 保留（原本存在）
        assert "momentum_20" in result.columns
        # 2. roe 删除（本次生成）
        assert "roe" not in result.columns
        # 3. 列与输入完全一致
        assert result.columns == data.columns

        # 验证筛选正确执行
        # momentum_20 > 0.05: 000001.SZ (0.1), 000003.SZ (0.08)
        # roe > 0.1: 000001.SZ (0.12), 000003.SZ (0.15)
        # 交集：000001.SZ, 000003.SZ
        assert len(result) == 2

        # 验证 _compute_factors 被调用（因为 roe 不存在）
        mock_compute.assert_called_once()

    def test_no_factor_computation_when_all_exist(self):
        """测试所有因子都已存在时，不调用 FactorEngine"""

        def filter_func(df: pl.DataFrame) -> pl.Series:
            return (df["factor_a"] > 0.5) & (df["factor_b"] < 0.3)

        manager = StockPoolManager(
            filter_func=filter_func,
            required_factors={
                "factor_a": "some_expr_a",
                "factor_b": "some_expr_b",
            },
        )

        # 输入数据已包含所有因子
        data = pl.DataFrame(
            {
                "ts_code": ["000001.SZ", "000002.SZ"],
                "trade_date": ["20240101"] * 2,
                "close": [10.0, 20.0],
                "factor_a": [0.6, 0.4],  # 原本存在
                "factor_b": [0.2, 0.5],  # 原本存在
            }
        )

        with patch("src.factors.engine.factor_engine.FactorEngine") as mock_engine:
            result = manager.filter_and_select_daily(data)

            # FactorEngine 不应被调用（所有因子都已存在）
            mock_engine.assert_not_called()

        # 验证结果正确
        assert len(result) == 1
        assert result["ts_code"][0] == "000001.SZ"
        assert result.columns == data.columns


class TestStockPoolManagerEdgeCases:
    """边界情况测试类"""

    def test_single_date(self):
        """测试单日数据"""

        def filter_func(df: pl.DataFrame) -> pl.Series:
            return df["close"] > 15

        manager = StockPoolManager(filter_func=filter_func)

        data = pl.DataFrame(
            {
                "ts_code": ["000001.SZ", "000002.SZ", "000003.SZ"],
                "trade_date": ["20240101"] * 3,
                "close": [10.0, 20.0, 30.0],
            }
        )

        result = manager.filter_and_select_daily(data)

        assert len(result) == 2
        assert result.columns == data.columns

    def test_single_stock_per_day(self):
        """测试每天只有一只股票"""

        def filter_func(df: pl.DataFrame) -> pl.Series:
            return df["close"] > 0  # 都保留

        manager = StockPoolManager(filter_func=filter_func)

        data = pl.DataFrame(
            {
                "ts_code": ["000001.SZ", "000002.SZ", "000003.SZ"],
                "trade_date": ["20240101", "20240102", "20240103"],
                "close": [10.0, 20.0, 30.0],
            }
        )

        result = manager.filter_and_select_daily(data)

        assert len(result) == 3
        assert result.columns == data.columns

    def test_filter_all_out_one_day(self):
        """测试某天全部过滤掉"""

        def filter_func(df: pl.DataFrame) -> pl.Series:
            return df["close"] > 100  # 很高的阈值

        manager = StockPoolManager(filter_func=filter_func)

        data = pl.DataFrame(
            {
                "ts_code": [
                    "000001.SZ",
                    "000002.SZ",
                    "000001.SZ",
                    "000002.SZ",
                ],
                "trade_date": [
                    "20240101",
                    "20240101",
                    "20240102",
                    "20240102",
                ],
                "close": [
                    10.0,
                    20.0,  # 日期1：都过滤掉
                    150.0,
                    200.0,  # 日期2：都保留
                ],
            }
        )

        result = manager.filter_and_select_daily(data)

        # 只有日期2的数据
        assert len(result) == 2
        assert all(result["trade_date"] == "20240102")

    def test_column_order_preserved(self):
        """测试列顺序保持不变"""

        def filter_func(df: pl.DataFrame) -> pl.Series:
            return df["close"] > 15

        manager = StockPoolManager(filter_func=filter_func)

        data = pl.DataFrame(
            {
                "ts_code": ["000001.SZ", "000002.SZ", "000003.SZ"],
                "trade_date": ["20240101"] * 3,
                "open": [9.0, 19.0, 29.0],
                "high": [11.0, 21.0, 31.0],
                "low": [8.0, 18.0, 28.0],
                "close": [10.0, 20.0, 30.0],
                "volume": [1000, 2000, 3000],
            }
        )

        result = manager.filter_and_select_daily(data)

        # 验证列顺序完全一致
        assert list(result.columns) == list(data.columns)


if __name__ == "__main__":
    pytest.main([__file__, "-v"])
