"""NaN 问题逐步诊断测试 - 精确定位问题环节

此测试会逐步检查 DataPipeline 的每个处理步骤，精确定位 NaN 产生的位置。
"""

import numpy as np
import polars as pl
import pytest

from src.factors import FactorEngine
from src.training import (
    FactorManager,
    DataPipeline,
    NullFiller,
    Winsorizer,
    StandardScaler,
)
from src.training.components.filters import STFilter
from src.experiment.common import (
    SELECTED_FACTORS,
    FACTOR_DEFINITIONS,
    LABEL_NAME,
    LABEL_FACTOR,
    stock_pool_filter,
    STOCK_FILTER_REQUIRED_COLUMNS,
)

# 测试配置
EXCLUDED_FACTORS = [f"GTJA_alpha{i:03d}" for i in range(1, 50)]  # 排除前50个加速测试
TEST_DATE_RANGE = {
    "train": ("20200101", "20201231"),  # 一整年数据
    "val": ("20210101", "20210331"),
    "test": ("20210401", "20210630"),
}


class TestNaNStepByStep:
    """逐步诊断 NaN 问题的测试类"""

    @pytest.fixture(scope="class")
    def base_data(self):
        """准备基础数据（未经过任何处理）"""
        print("\n" + "=" * 80)
        print("[Fixture] 准备基础数据...")

        engine = FactorEngine()
        factor_manager = FactorManager(
            selected_factors=SELECTED_FACTORS,
            factor_definitions=FACTOR_DEFINITIONS,
            label_factor=LABEL_FACTOR,
            excluded_factors=EXCLUDED_FACTORS,
        )

        # 注册因子
        feature_cols = factor_manager.register_to_engine(engine, verbose=False)
        print(f"  特征数: {len(feature_cols)}")

        # 计算完整日期范围
        all_start = min(
            TEST_DATE_RANGE["train"][0],
            TEST_DATE_RANGE["val"][0],
            TEST_DATE_RANGE["test"][0],
        )
        all_end = max(
            TEST_DATE_RANGE["train"][1],
            TEST_DATE_RANGE["val"][1],
            TEST_DATE_RANGE["test"][1],
        )

        # 计算因子
        raw_data = engine.compute(
            factor_names=feature_cols + [LABEL_NAME],
            start_date=all_start,
            end_date=all_end,
        )
        print(f"  原始数据形状: {raw_data.shape}")

        return {
            "engine": engine,
            "factor_manager": factor_manager,
            "feature_cols": feature_cols,
            "raw_data": raw_data,
        }

    def test_step_0_raw_data(self, base_data):
        """步骤0: 检查原始数据中的 NaN"""
        print("\n" + "=" * 80)
        print("[步骤0] 检查原始数据中的 NaN")

        raw_data = base_data["raw_data"]
        feature_cols = base_data["feature_cols"]

        nan_stats = self._check_nan_in_df(raw_data, feature_cols, "原始数据")

        # 记录有 NaN 的列
        print(f"  含 NaN 的特征列数: {len(nan_stats['cols_with_nan'])}")
        if nan_stats["cols_with_nan"]:
            print(f"  示例: {nan_stats['cols_with_nan'][:5]}")

        return nan_stats

    def test_step_1_after_st_filter(self, base_data):
        """步骤1: 检查 STFilter 后的 NaN"""
        print("\n" + "=" * 80)
        print("[步骤1] 检查 STFilter 后的 NaN")

        raw_data = base_data["raw_data"]
        feature_cols = base_data["feature_cols"]
        engine = base_data["engine"]

        st_filter = STFilter(data_router=engine.router)
        filtered_data = st_filter.filter(raw_data)

        print(f"  过滤后数据形状: {filtered_data.shape}")
        print(f"  删除记录数: {len(raw_data) - len(filtered_data)}")

        nan_stats = self._check_nan_in_df(filtered_data, feature_cols, "STFilter后")

        # 对比步骤0，看是否有新增 NaN
        step0_nan = self.test_step_0_raw_data(base_data)
        if nan_stats["total_nan"] != step0_nan["total_nan"]:
            print(
                f"  [警告] NaN 数量变化: {step0_nan['total_nan']} -> {nan_stats['total_nan']}"
            )

        return nan_stats

    def test_step_2_after_stock_pool(self, base_data):
        """步骤2: 检查股票池筛选后的 NaN"""
        print("\n" + "=" * 80)
        print("[步骤2] 检查股票池筛选后的 NaN")

        raw_data = base_data["raw_data"]
        feature_cols = base_data["feature_cols"]
        engine = base_data["engine"]

        # 先应用 STFilter
        st_filter = STFilter(data_router=engine.router)
        filtered_data = st_filter.filter(raw_data)

        # 再应用股票池筛选
        from src.training.core.stock_pool_manager import StockPoolManager

        pool_manager = StockPoolManager(
            filter_func=stock_pool_filter,
            required_columns=STOCK_FILTER_REQUIRED_COLUMNS,
            data_router=engine.router,
        )
        pool_data = pool_manager.filter_and_select_daily(filtered_data)

        print(f"  筛选后数据形状: {pool_data.shape}")
        print(f"  删除记录数: {len(filtered_data) - len(pool_data)}")

        nan_stats = self._check_nan_in_df(pool_data, feature_cols, "股票池筛选后")
        return nan_stats

    def test_step_3_train_split_without_skip(self, base_data):
        """步骤3: 检查训练集划分后的 NaN（不跳过天数）"""
        print("\n" + "=" * 80)
        print("[步骤3] 检查训练集划分后的 NaN（不跳过天数）")

        raw_data = base_data["raw_data"]
        feature_cols = base_data["feature_cols"]
        engine = base_data["engine"]

        # 应用过滤器
        st_filter = STFilter(data_router=engine.router)
        filtered_data = st_filter.filter(raw_data)

        from src.training.core.stock_pool_manager import StockPoolManager

        pool_manager = StockPoolManager(
            filter_func=stock_pool_filter,
            required_columns=STOCK_FILTER_REQUIRED_COLUMNS,
            data_router=engine.router,
        )
        pool_data = pool_manager.filter_and_select_daily(filtered_data)

        # 划分训练集
        train_start, train_end = TEST_DATE_RANGE["train"]
        train_mask = (pool_data["trade_date"] >= train_start) & (
            pool_data["trade_date"] <= train_end
        )
        train_df = pool_data.filter(train_mask)

        print(f"  训练集形状: {train_df.shape}")

        # 统计交易日数量
        unique_dates = train_df["trade_date"].unique().sort()
        print(f"  训练集交易日数量: {len(unique_dates)}")
        print(f"  日期范围: {unique_dates[0]} ~ {unique_dates[-1]}")

        nan_stats = self._check_nan_in_df(train_df, feature_cols, "训练集(不跳过)")

        # 返回训练集供后续测试使用
        return {
            "nan_stats": nan_stats,
            "train_df": train_df,
            "unique_dates": unique_dates,
        }

    def test_step_4_train_split_with_skip(self, base_data):
        """步骤4: 检查训练集划分后的 NaN（跳过前252天）"""
        print("\n" + "=" * 80)
        print("[步骤4] 检查训练集划分后的 NaN（跳过前252天）")

        step3_result = self.test_step_3_train_split_without_skip(base_data)
        train_df = step3_result["train_df"]
        unique_dates = step3_result["unique_dates"]
        feature_cols = base_data["feature_cols"]

        # 跳过前252天
        skip_days = 252
        if len(unique_dates) > skip_days:
            start_date = unique_dates[skip_days]
            train_df_skipped = train_df.filter(pl.col("trade_date") >= start_date)
            print(f"  跳过前{skip_days}天后，从 {start_date} 开始")
            print(f"  跳过后训练集形状: {train_df_skipped.shape}")
            print(f"  跳过记录数: {len(train_df) - len(train_df_skipped)}")
        else:
            train_df_skipped = train_df
            print(
                f"  [警告] 训练集交易日数({len(unique_dates)})少于跳过天数({skip_days})，未跳过"
            )

        nan_stats = self._check_nan_in_df(
            train_df_skipped, feature_cols, "训练集(跳过252天)"
        )

        return {
            "nan_stats": nan_stats,
            "train_df": train_df_skipped,
        }

    def test_step_5_after_null_filler(self, base_data):
        """步骤5: 检查 NullFiller 后的 NaN"""
        print("\n" + "=" * 80)
        print("[步骤5] 检查 NullFiller 后的 NaN")

        step4_result = self.test_step_4_train_split_with_skip(base_data)
        train_df = step4_result["train_df"]
        feature_cols = base_data["feature_cols"]

        print(f"  处理前数据形状: {train_df.shape}")

        # 应用 NullFiller
        null_filler = NullFiller(
            feature_cols=feature_cols, strategy="mean", by_date=True
        )
        after_null = null_filler.fit_transform(train_df)

        print(f"  处理后数据形状: {after_null.shape}")

        nan_stats = self._check_nan_in_df(after_null, feature_cols, "NullFiller后")

        # 检查哪些列还有 NaN
        if nan_stats["cols_with_nan"]:
            print(
                f"  [错误] NullFiller 后仍有 {len(nan_stats['cols_with_nan'])} 列含 NaN:"
            )
            for col in nan_stats["cols_with_nan"][:10]:
                count = after_null[col].null_count()
                dtype = after_null[col].dtype
                print(f"    {col}: {count} 个 NaN, dtype={dtype}")

        return {
            "nan_stats": nan_stats,
            "after_null": after_null,
        }

    def test_step_6_after_winsorizer(self, base_data):
        """步骤6: 检查 Winsorizer 后的 NaN"""
        print("\n" + "=" * 80)
        print("[步骤6] 检查 Winsorizer 后的 NaN")

        step5_result = self.test_step_5_after_null_filler(base_data)
        after_null = step5_result["after_null"]
        feature_cols = base_data["feature_cols"]

        # 应用 Winsorizer
        winsorizer = Winsorizer(
            feature_cols=feature_cols, lower=0.01, upper=0.99, by_date=False
        )
        after_winsor = winsorizer.fit_transform(after_null)

        nan_stats = self._check_nan_in_df(after_winsor, feature_cols, "Winsorizer后")

        # 检查哪些列还有 NaN
        if nan_stats["cols_with_nan"]:
            print(
                f"  [错误] Winsorizer 后仍有 {len(nan_stats['cols_with_nan'])} 列含 NaN:"
            )
            for col in nan_stats["cols_with_nan"][:10]:
                count = after_winsor[col].null_count()
                dtype = after_winsor[col].dtype
                print(f"    {col}: {count} 个 NaN, dtype={dtype}")

        return {
            "nan_stats": nan_stats,
            "after_winsor": after_winsor,
        }

    def test_step_7_after_standard_scaler(self, base_data):
        """步骤7: 检查 StandardScaler 后的 NaN"""
        print("\n" + "=" * 80)
        print("[步骤7] 检查 StandardScaler 后的 NaN")

        step6_result = self.test_step_6_after_winsorizer(base_data)
        after_winsor = step6_result["after_winsor"]
        feature_cols = base_data["feature_cols"]

        # 在应用 StandardScaler 之前，检查那些后来出问题的列
        print("\n  [预检查] StandardScaler 前，检查关键列...")
        problematic_cols = [
            "GTJA_alpha062",
            "GTJA_alpha073",
            "GTJA_alpha085",
            "GTJA_alpha087",
            "GTJA_alpha092",
            "GTJA_alpha103",
            "GTJA_alpha104",
            "GTJA_alpha117",
            "GTJA_alpha124",
            "GTJA_alpha131",
        ]
        for col in problematic_cols:
            if col in after_winsor.columns:
                null_count = after_winsor[col].null_count()
                dtype = after_winsor[col].dtype
                min_val = after_winsor[col].min()
                max_val = after_winsor[col].max()
                print(
                    f"    {col}: null={null_count}, dtype={dtype}, min={min_val}, max={max_val}"
                )

        # 应用 StandardScaler
        scaler = StandardScaler(feature_cols=feature_cols)
        after_scaler = scaler.fit_transform(after_winsor)

        # 检查 StandardScaler 学到的统计量
        print("\n  [统计量检查] StandardScaler 学到的统计量...")
        for col in problematic_cols:
            if col in scaler.mean_:
                print(f"    {col}: mean={scaler.mean_[col]}, std={scaler.std_[col]}")
            else:
                print(f"    {col}: [未学到统计量]")

        nan_stats = self._check_nan_in_df(
            after_scaler, feature_cols, "StandardScaler后"
        )

        # 检查哪些列还有 NaN
        if nan_stats["cols_with_nan"]:
            print(
                f"  [错误] StandardScaler 后仍有 {len(nan_stats['cols_with_nan'])} 列含 NaN:"
            )
            for col in nan_stats["cols_with_nan"][:10]:
                count = after_scaler[col].null_count()
                dtype = after_scaler[col].dtype
                print(f"    {col}: {count} 个 NaN, dtype={dtype}")

                # 检查这列是否学到了统计量
                if col in scaler.mean_:
                    print(
                        f"      mean={scaler.mean_[col]:.4f}, std={scaler.std_[col]:.4f}"
                    )
                else:
                    print(f"      [警告] 未学到统计量！")

        return {
            "nan_stats": nan_stats,
            "after_scaler": after_scaler,
            "scaler": scaler,
        }

    def test_step_8_extract_X(self, base_data):
        """步骤8: 检查提取 X 后的 NaN（转换为 numpy）"""
        print("\n" + "=" * 80)
        print("[步骤8] 检查提取 X 后的 NaN")

        step7_result = self.test_step_7_after_standard_scaler(base_data)
        after_scaler = step7_result["after_scaler"]
        feature_cols = base_data["feature_cols"]

        # 提取 X
        X_df = after_scaler.select(feature_cols)
        print(f"  X DataFrame 形状: {X_df.shape}")

        # 对比 DataFrame 和 select 后的 null 数量
        print("\n  [对比] DataFrame vs select 后的 null 数量:")
        mismatched = []
        for col in feature_cols[:20]:  # 只检查前20个
            null_in_df = after_scaler[col].null_count()
            null_in_x = X_df[col].null_count()
            if null_in_df != null_in_x:
                mismatched.append((col, null_in_df, null_in_x))

        if mismatched:
            print(f"    [警告] 发现 {len(mismatched)} 列不匹配:")
            for col, df_null, x_null in mismatched[:10]:
                print(f"      {col}: DataFrame={df_null}, X={x_null}")
        else:
            print(f"    [通过] 所有列的 null 数量一致")

        # 转换为 numpy
        X_np = X_df.to_numpy()
        print(f"\n  X numpy 形状: {X_np.shape}")

        nan_count = np.isnan(X_np).sum()
        print(f"  X 中 NaN 总数: {nan_count}")

        if nan_count > 0:
            # 找出哪些列有 NaN
            nan_by_col = []
            for i, col in enumerate(feature_cols):
                col_nan = np.isnan(X_np[:, i]).sum()
                if col_nan > 0:
                    nan_by_col.append((col, col_nan))

            print(f"\n  [错误] 含 NaN 的特征列数: {len(nan_by_col)}")
            for col, count in nan_by_col[:10]:
                # 检查原始 DataFrame 中的情况
                df_null = after_scaler[col].null_count()
                dtype = after_scaler[col].dtype

                # 检查是否有 Infinity
                inf_count_pos = (after_scaler[col] == float("inf")).sum()
                inf_count_neg = (after_scaler[col] == float("-inf")).sum()

                # 检查 min/max
                col_min = after_scaler[col].min()
                col_max = after_scaler[col].max()

                print(
                    f"    {col}: numpy中{count}个NaN, DataFrame中{df_null}个null, dtype={dtype}"
                )
                print(f"      min={col_min}, max={col_max}")
                print(f"      +inf={inf_count_pos}, -inf={inf_count_neg}")

                # 如果有 inf，显示一些样本值
                if inf_count_pos > 0 or inf_count_neg > 0:
                    sample_vals = after_scaler[col].drop_nulls().tail(5).to_list()
                    print(f"      样本值: {sample_vals}")

            # 断言失败，显示详细信息
            assert False, f"X 中含 {nan_count} 个 NaN，涉及 {len(nan_by_col)} 个特征列"
        else:
            print("\n  [通过] X 中无 NaN！")

    def _check_nan_in_df(
        self, df: pl.DataFrame, feature_cols: list, step_name: str
    ) -> dict:
        """检查 DataFrame 中的 NaN 统计信息

        Returns:
            dict: {
                'total_nan': 总NaN数,
                'cols_with_nan': 含NaN的列名列表,
                'nan_by_col': {列名: NaN数} 的字典
            }
        """
        nan_by_col = {}
        total_nan = 0

        for col in feature_cols:
            null_count = df[col].null_count()
            if null_count > 0:
                nan_by_col[col] = null_count
                total_nan += null_count

        cols_with_nan = list(nan_by_col.keys())

        print(f"  {step_name}:")
        print(f"    总记录数: {len(df)}")
        print(f"    特征列数: {len(feature_cols)}")
        print(f"    总NaN数: {total_nan}")
        print(f"    含NaN的列数: {len(cols_with_nan)}")

        if cols_with_nan and len(cols_with_nan) <= 5:
            print(f"    含NaN的列: {cols_with_nan}")
        elif cols_with_nan:
            print(f"    含NaN的列(前5): {cols_with_nan[:5]}...")

        return {
            "total_nan": total_nan,
            "cols_with_nan": cols_with_nan,
            "nan_by_col": nan_by_col,
        }


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