import numpy as np
import polars as pl
import pytest

from src.factorminer.evaluation.stock_pool_registry import StockPoolRegistry


class TestStockPoolRegistry:
    def test_add_and_get_pool_names(self):
        registry = StockPoolRegistry()
        registry.add_pool("all", lambda df: pl.Series([True] * len(df)))
        registry.add_pool("growth", lambda df: df["ts_code"].str.starts_with("300"))
        assert registry.get_pool_names() == ["all", "growth"]

    def test_build_masks(self):
        registry = StockPoolRegistry()
        registry.add_pool("all", lambda df: pl.Series([True] * len(df)))
        registry.add_pool("growth", lambda df: df["ts_code"].str.starts_with("300"))
        asset_codes = ["000001.SZ", "300001.SZ", "688001.SH"]
        masks = registry.build_masks(asset_codes)
        assert masks["all"].sum() == 3
        assert masks["growth"].sum() == 1
        assert masks["growth"][1]

    def test_filter_signals(self):
        registry = StockPoolRegistry()
        registry.add_pool("growth", lambda df: df["ts_code"].str.starts_with("300"))
        asset_codes = ["000001.SZ", "300001.SZ", "688001.SH"]
        registry.build_masks(asset_codes)
        signals = np.array([[1, 2], [3, 4], [5, 6]], dtype=np.float64)
        filtered = registry.filter_signals(signals, "growth")
        assert filtered.shape == (1, 2)
        np.testing.assert_array_equal(filtered, [[3, 4]])

    def test_build_masks_with_metadata(self):
        registry = StockPoolRegistry()
        registry.add_pool(
            "small_cap",
            lambda df: df["ts_code"].is_in(
                df.sort("total_mv").head(2)["ts_code"].implode()
            ),
            required_columns=["total_mv"],
        )
        asset_codes = ["A", "B", "C"]
        metadata = pl.DataFrame(
            {
                "ts_code": ["A", "B", "C"],
                "total_mv": [300.0, 100.0, 200.0],
            }
        )
        masks = registry.build_masks(asset_codes, metadata_df=metadata)
        assert masks["small_cap"].sum() == 2
        assert masks["small_cap"][1] and masks["small_cap"][2]

    def test_expr_fallback(self):
        registry = StockPoolRegistry()
        registry.add_pool(
            "expr_pool", lambda df: pl.col("ts_code").str.starts_with("300")
        )
        asset_codes = ["000001.SZ", "300001.SZ"]
        masks = registry.build_masks(asset_codes)
        assert masks["expr_pool"].sum() == 1

    def test_pre_filters(self):
        registry = StockPoolRegistry(
            pre_filters=[lambda df: df["ts_code"].str.starts_with("0")]
        )
        registry.add_pool("all", lambda df: pl.Series([True] * len(df)))
        asset_codes = ["000001.SZ", "300001.SZ", "688001.SH"]
        masks = registry.build_masks(asset_codes)
        assert masks["all"].sum() == 1
        assert masks["all"][0]
