"""TabM NaN 问题诊断测试 + train_skip_days 功能验证

诊断 loss 为 nan 的根因：
1. 标签中是否有 NaN 或极端值
2. 标准化后是否有 NaN
3. 是否有方差为0或接近0的列
4. train_skip_days 功能是否正常工作
"""

import numpy as np
import polars as pl
import pytest

from src.factors import FactorEngine
from src.training import (
    FactorManager,
    DataPipeline,
    TabMRegressionTask,
    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,
)


# TabM 模型参数（简化用于快速测试）
MODEL_PARAMS = {
    "n_blocks": 2,
    "d_block": 128,
    "dropout": 0.1,
    "ensemble_size": 4,  # 简化
    "batch_size": 256,
    "learning_rate": 1e-4,  # 降低学习率
    "weight_decay": 1e-5,
    "epochs": 3,
    "early_stopping_patience": 5,
}

# 测试日期范围
TEST_DATE_RANGE = {
    "train": ("20200101", "20200630"),
    "val": ("20200701", "20200731"),
    "test": ("20200801", "20200831"),
}

# 小范围测试用的因子排除列表
EXCLUDED_FACTORS = [
    "GTJA_alpha001",
    "GTJA_alpha002",
    "GTJA_alpha003",
    "GTJA_alpha004",
    "GTJA_alpha005",
    "GTJA_alpha006",
    "GTJA_alpha007",
    "GTJA_alpha008",
    "GTJA_alpha009",
    "GTJA_alpha010",
]


class TestTabMNanDebug:
    """TabM NaN 问题诊断测试类（使用 DataPipeline）"""

    @pytest.fixture(scope="class")
    def engine_and_factor_manager(self):
        """准备 FactorEngine 和 FactorManager"""
        engine = FactorEngine()

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

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

    @pytest.fixture(scope="class")
    def pipeline_data(self, engine_and_factor_manager):
        """使用 DataPipeline 准备数据，并过滤掉全为NaN的列"""
        engine = engine_and_factor_manager["engine"]
        factor_manager = engine_and_factor_manager["factor_manager"]

        # 创建 DataPipeline（使用 train_skip_days=0 进行基础测试）
        pipeline = DataPipeline(
            factor_manager=factor_manager,
            processor_configs=[
                (NullFiller, {"strategy": "mean"}),
                (Winsorizer, {"lower": 0.01, "upper": 0.99}),
                (StandardScaler, {}),
            ],
            filters=[STFilter(data_router=engine.router)],
            stock_pool_filter_func=stock_pool_filter,
            stock_pool_required_columns=STOCK_FILTER_REQUIRED_COLUMNS,
            train_skip_days=0,  # 先不跳过天数进行测试
        )

        # 准备数据
        data = pipeline.prepare_data(
            engine=engine,
            date_range=TEST_DATE_RANGE,
            label_name=LABEL_NAME,
            verbose=True,
        )

        # 获取特征列
        feature_cols = factor_manager.register_to_engine(engine, verbose=False)

        # 过滤掉全为NaN的列（这些因子计算失败，没有修复意义）
        print("\n[DataPipeline] 检查并过滤全为NaN的特征列...")
        for split_name in ["train", "val", "test"]:
            X_df = data[split_name]["X"]
            y_series = data[split_name]["y"]
            raw_data = data[split_name]["raw_data"]

            # 删除标签为NaN的行
            y_nan_count = y_series.null_count()
            if y_nan_count > 0:
                print(f"  {split_name}: 发现 {y_nan_count} 个标签为NaN的行，将被删除")
                # 创建有效标签的mask
                valid_mask = y_series.is_not_null()
                # 过滤所有相关数据
                X_df = X_df.filter(valid_mask)
                y_series = y_series.filter(valid_mask)
                raw_data = raw_data.filter(valid_mask)
                # 更新数据
                data[split_name]["X"] = X_df
                data[split_name]["y"] = y_series
                data[split_name]["raw_data"] = raw_data

            # 检查每列的NaN数量
            nan_counts = {col: X_df[col].null_count() for col in X_df.columns}
            total_rows = len(X_df)

            # 找出全为NaN的列
            all_nan_cols = [
                col for col, count in nan_counts.items() if count == total_rows
            ]

            if all_nan_cols:
                print(
                    f"  {split_name}: 发现 {len(all_nan_cols)} 个全为NaN的列，将被删除"
                )
                print(
                    f"    列名: {all_nan_cols[:5]}{'...' if len(all_nan_cols) > 5 else ''}"
                )

                # 更新feature_cols
                feature_cols = [c for c in feature_cols if c not in all_nan_cols]

                # 从X中删除这些列
                X_df = X_df.select(feature_cols)

                # 从raw_data中也删除这些列（保留原始特征列以外的列如trade_date, ts_code等）
                raw_cols_to_keep = [
                    c for c in raw_data.columns if c not in all_nan_cols
                ]
                raw_data = raw_data.select(raw_cols_to_keep)

                # 更新数据
                data[split_name]["X"] = X_df
                data[split_name]["raw_data"] = raw_data
                data[split_name]["feature_cols"] = feature_cols

            # 验证没有NaN了
            X_np = X_df.to_numpy()
            nan_count = np.isnan(X_np).sum()
            assert nan_count == 0, f"{split_name} 中仍有 {nan_count} 个NaN"

        print(f"  过滤后特征数: {len(feature_cols)}")
        print("  [通过] 所有特征列均无NaN")

        return {
            "pipeline": pipeline,
            "data": data,
            "feature_cols": feature_cols,
            "engine": engine,
        }

    def test_pipeline_data_structure(self, pipeline_data):
        """诊断0: DataPipeline 返回的数据结构检查"""
        data = pipeline_data["data"]
        feature_cols = pipeline_data["feature_cols"]

        print("\n[诊断0] DataPipeline 数据结构检查:")

        # 检查数据结构
        assert "train" in data, "缺少 train 数据"
        assert "val" in data, "缺少 val 数据"
        assert "test" in data, "缺少 test 数据"

        for split_name in ["train", "val", "test"]:
            split_data = data[split_name]
            assert "X" in split_data, f"{split_name} 缺少 X"
            assert "y" in split_data, f"{split_name} 缺少 y"
            assert "raw_data" in split_data, f"{split_name} 缺少 raw_data"
            assert "feature_cols" in split_data, f"{split_name} 缺少 feature_cols"

            X = split_data["X"]
            y = split_data["y"]

            print(f"  {split_name}:")
            print(f"    X 形状: {X.shape}")
            print(f"    y 形状: {len(y)}")
            print(f"    特征数: {len(split_data['feature_cols'])}")

            # 检查维度一致性
            assert X.shape[0] == len(y), f"{split_name} X 和 y 行数不匹配"
            assert X.shape[1] == len(feature_cols), f"{split_name} 特征数不匹配"

        print("  [通过] 数据结构正确")

    def test_label_quality_with_pipeline(self, pipeline_data):
        """诊断1: 使用 DataPipeline 后的标签数据质量检查"""
        data = pipeline_data["data"]

        # 检查所有数据集的标签质量
        for split_name in ["train", "val", "test"]:
            y = data[split_name]["y"]
            y_np = y.to_numpy()

            print(f"\n[诊断1-{split_name}] 标签数据质量:")
            print(f"  总数: {len(y)}")
            print(f"  NaN数量: {y.null_count()}")
            print(f"  均值: {y.mean():.6f}")
            print(f"  标准差: {y.std():.6f}")
            print(f"  最小值: {y.min():.6f}")
            print(f"  最大值: {y.max():.6f}")

            inf_count = np.isinf(y_np).sum()
            nan_count = np.isnan(y_np).sum()
            print(f"  inf数量: {inf_count}")
            print(f"  nan数量: {nan_count}")

            assert inf_count == 0, f"{split_name} 标签含 inf: {inf_count}"
            assert nan_count == 0, f"{split_name} 标签含 nan: {nan_count}"

    def test_processed_data_quality_with_pipeline(self, pipeline_data):
        """诊断2: 使用 DataPipeline 处理后的特征数据质量"""
        data = pipeline_data["data"]
        feature_cols = pipeline_data["feature_cols"]

        for split_name in ["train", "val", "test"]:
            X = data[split_name]["X"]
            y = data[split_name]["y"]

            X_np = X.to_numpy().astype(np.float32)
            y_np = y.to_numpy().astype(np.float32)

            print(f"\n[诊断2-{split_name}] 处理后数据质量:")
            print(f"  X 形状: {X_np.shape}, dtype: {X_np.dtype}")
            print(f"  y 形状: {y_np.shape}, dtype: {y_np.dtype}")
            print(f"  X中NaN: {np.isnan(X_np).sum()}")
            print(f"  X中Inf: {np.isinf(X_np).sum()}")
            print(f"  y中NaN: {np.isnan(y_np).sum()}")
            print(f"  y中Inf: {np.isinf(y_np).sum()}")

            assert np.isnan(X_np).sum() == 0, f"{split_name} X含NaN"
            assert np.isnan(y_np).sum() == 0, f"{split_name} y含NaN"
            assert np.isinf(X_np).sum() == 0, f"{split_name} X含Inf"
            assert np.isinf(y_np).sum() == 0, f"{split_name} y含Inf"

    def test_train_skip_days_functionality(self, engine_and_factor_manager):
        """诊断3: train_skip_days 功能验证"""
        engine = engine_and_factor_manager["engine"]
        factor_manager = engine_and_factor_manager["factor_manager"]

        print("\n[诊断3] train_skip_days 功能验证:")

        # 创建一个跳过50天数据的 pipeline
        skip_days = 50
        pipeline_with_skip = DataPipeline(
            factor_manager=factor_manager,
            processor_configs=[
                (NullFiller, {"strategy": "mean"}),
                (Winsorizer, {"lower": 0.01, "upper": 0.99}),
                (StandardScaler, {}),
            ],
            filters=[STFilter(data_router=engine.router)],
            stock_pool_filter_func=stock_pool_filter,
            stock_pool_required_columns=STOCK_FILTER_REQUIRED_COLUMNS,
            train_skip_days=skip_days,
        )

        # 准备数据
        data_with_skip = pipeline_with_skip.prepare_data(
            engine=engine,
            date_range=TEST_DATE_RANGE,
            label_name=LABEL_NAME,
            verbose=True,
        )

        # 获取原始数据（不跳过天数）用于对比
        pipeline_no_skip = DataPipeline(
            factor_manager=factor_manager,
            processor_configs=[
                (NullFiller, {"strategy": "mean"}),
                (Winsorizer, {"lower": 0.01, "upper": 0.99}),
                (StandardScaler, {}),
            ],
            filters=[STFilter(data_router=engine.router)],
            stock_pool_filter_func=stock_pool_filter,
            stock_pool_required_columns=STOCK_FILTER_REQUIRED_COLUMNS,
            train_skip_days=0,
        )

        data_no_skip = pipeline_no_skip.prepare_data(
            engine=engine,
            date_range=TEST_DATE_RANGE,
            label_name=LABEL_NAME,
            verbose=False,
        )

        # 验证训练数据减少了
        train_with_skip = data_with_skip["train"]["raw_data"]
        train_no_skip = data_no_skip["train"]["raw_data"]

        print(f"\n  对比结果:")
        print(f"    不跳过时训练数据: {len(train_no_skip)} 条")
        print(f"    跳过{skip_days}天后: {len(train_with_skip)} 条")
        print(f"    减少: {len(train_no_skip) - len(train_with_skip)} 条")

        # 验证验证集和测试集不受影响
        assert len(data_with_skip["val"]["raw_data"]) == len(
            data_no_skip["val"]["raw_data"]
        ), "val 数据不应该受 train_skip_days 影响"
        assert len(data_with_skip["test"]["raw_data"]) == len(
            data_no_skip["test"]["raw_data"]
        ), "test 数据不应该受 train_skip_days 影响"

        # 验证日期确实被跳过了
        if len(train_no_skip) > 0:
            dates_no_skip = sorted(train_no_skip["trade_date"].unique())
            dates_with_skip = sorted(train_with_skip["trade_date"].unique())

            print(f"\n  日期对比:")
            print(
                f"    不跳过 - 最早日期: {dates_no_skip[0]}, 共 {len(dates_no_skip)} 个交易日"
            )
            print(
                f"    跳过 - 最早日期: {dates_with_skip[0]}, 共 {len(dates_with_skip)} 个交易日"
            )

            # 验证跳过的数据确实从更晚的日期开始
            if len(dates_no_skip) > skip_days:
                expected_start_date = dates_no_skip[skip_days]
                assert dates_with_skip[0] == expected_start_date, (
                    f"预期从 {expected_start_date} 开始，实际从 {dates_with_skip[0]} 开始"
                )
                print(f"    [通过] 正确跳过前 {skip_days} 个交易日")

    def test_training_with_pipeline(self, pipeline_data):
        """诊断4: 使用 DataPipeline 处理后的数据进行训练测试"""
        import torch
        from torch.utils.data import DataLoader, TensorDataset
        from tabm import TabM

        data = pipeline_data["data"]
        feature_cols = pipeline_data["feature_cols"]

        # 获取训练数据
        X_train_df = data["train"]["X"]
        y_train_series = data["train"]["y"]

        # 删除标签为NaN的行
        valid_mask = y_train_series.is_not_null()
        y_nan_count = y_train_series.null_count()
        if y_nan_count > 0:
            print(f"  发现 {y_nan_count} 个标签为NaN的行，将被删除")
            X_train_df = X_train_df.filter(valid_mask)
            y_train_series = y_train_series.filter(valid_mask)

        X_train = X_train_df.to_numpy().astype(np.float32)
        y_train = y_train_series.to_numpy().astype(np.float32)

        # 只取前1000条加速测试
        X_train = X_train[:1000]
        y_train = y_train[:1000]

        print(f"\n[诊断4] DataPipeline 数据训练测试:")
        print(f"  X 形状: {X_train.shape}")
        print(f"  y 形状: {y_train.shape}")
        print(f"  X中NaN: {np.isnan(X_train).sum()}, Inf: {np.isinf(X_train).sum()}")
        print(f"  y中NaN: {np.isnan(y_train).sum()}, Inf: {np.isinf(y_train).sum()}")

        # 创建 TabM 模型
        n_features = X_train.shape[1]
        model = TabM.make(
            n_num_features=n_features,
            cat_cardinalities=[],
            d_out=1,
            n_blocks=2,
            d_block=128,
            dropout=0.1,
            k=4,
        )

        # 训练一个 epoch
        dataset = TensorDataset(torch.from_numpy(X_train), torch.from_numpy(y_train))
        loader = DataLoader(dataset, batch_size=256, shuffle=True)

        optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
        criterion = torch.nn.MSELoss()

        model.train()
        losses = []

        for batch_idx, (bx, by) in enumerate(loader):
            optimizer.zero_grad()
            outputs = model(bx)  # [B, E, 1]
            outputs_squeezed = outputs.squeeze(-1)  # [B, E]
            by_expanded = by.unsqueeze(-1).expand(-1, 4)  # [B, E]
            loss = criterion(outputs_squeezed, by_expanded)
            loss.backward()
            optimizer.step()

            losses.append(loss.item())
            print(f"  批次 {batch_idx + 1} loss: {loss.item():.6f}")

            assert not torch.isnan(loss), f"批次 {batch_idx + 1} loss 为 nan!"

            if batch_idx >= 2:  # 只测试前3个批次
                break

        print(f"  [通过] 所有批次 loss 正常，无 NaN")


class TestTrainSkipDaysEdgeCases:
    """train_skip_days 边界情况测试"""

    @pytest.fixture(scope="class")
    def engine_and_factor_manager(self):
        """准备 FactorEngine 和 FactorManager"""
        engine = FactorEngine()

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

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

    def test_skip_more_than_available_days(self, engine_and_factor_manager):
        """测试跳过天数超过可用天数的情况"""
        engine = engine_and_factor_manager["engine"]
        factor_manager = engine_and_factor_manager["factor_manager"]

        print("\n[边界测试] 跳过天数超过可用天数:")

        # 使用一个很大的跳过天数
        pipeline = DataPipeline(
            factor_manager=factor_manager,
            processor_configs=[(NullFiller, {"strategy": "mean"})],
            filters=[STFilter(data_router=engine.router)],
            stock_pool_filter_func=stock_pool_filter,
            stock_pool_required_columns=STOCK_FILTER_REQUIRED_COLUMNS,
            train_skip_days=1000,  # 超过测试期间的交易日数
        )

        # 准备数据（应该能正常运行并发出警告）
        data = pipeline.prepare_data(
            engine=engine,
            date_range=TEST_DATE_RANGE,
            label_name=LABEL_NAME,
            verbose=True,
        )

        # 即使跳过天数很多，也应该有数据返回
        # 如果交易日数少于跳过天数，应该保留所有数据（只发出警告）
        train_data = data["train"]["raw_data"]
        print(f"  训练数据量: {len(train_data)} 条")
        print(f"  [通过] 程序未崩溃")

    def test_skip_zero_days(self, engine_and_factor_manager):
        """测试跳过0天（即不跳过）"""
        engine = engine_and_factor_manager["engine"]
        factor_manager = engine_and_factor_manager["factor_manager"]

        print("\n[边界测试] 跳过0天:")

        pipeline = DataPipeline(
            factor_manager=factor_manager,
            processor_configs=[(NullFiller, {"strategy": "mean"})],
            filters=[STFilter(data_router=engine.router)],
            stock_pool_filter_func=stock_pool_filter,
            stock_pool_required_columns=STOCK_FILTER_REQUIRED_COLUMNS,
            train_skip_days=0,
        )

        data = pipeline.prepare_data(
            engine=engine,
            date_range=TEST_DATE_RANGE,
            label_name=LABEL_NAME,
            verbose=False,
        )

        train_data = data["train"]["raw_data"]
        print(f"  训练数据量: {len(train_data)} 条")
        print(f"  [通过] skip=0 时数据正常")


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