"""探针法因子筛选 - 真实数据集成

使用真实因子数据和训练流程执行探针筛选。

使用方法:
    uv run python src/experiment/probe_selection/run_probe_selection.py
"""

import os
import sys
from typing import List

import polars as pl

# 添加项目根目录到路径
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..", ".."))

from src.experiment.probe_selection import ProbeSelector
from src.factors import FactorEngine
from src.training import (
    DateSplitter,
    NullFiller,
    StandardScaler,
    StockPoolManager,
    STFilter,
    Winsorizer,
    check_data_quality,
)
from src.training.components.models.lightgbm import LightGBMModel


# 配置参数
LABEL_NAME = "future_return_5"

# 完整因子列表（来自 regression.py）
SELECTED_FACTORS = [
    # ================= 1. 价格、趋势与路径依赖 =================
    "ma_5",
    "ma_20",
    "ma_ratio_5_20",
    "bias_10",
    "high_low_ratio",
    "bbi_ratio",
    "return_5",
    "return_20",
    "kaufman_ER_20",
    "mom_acceleration_10_20",
    "drawdown_from_high_60",
    "up_days_ratio_20",
    # ================= 2. 波动率、风险调整与高阶矩 =================
    "volatility_5",
    "volatility_20",
    "volatility_ratio",
    "std_return_20",
    "sharpe_ratio_20",
    "min_ret_20",
    "volatility_squeeze_5_60",
    # ================= 3. 日内微观结构与异象 =================
    "overnight_intraday_diff",
    "upper_shadow_ratio",
    "capital_retention_20",
    "max_ret_20",
    # ================= 4. 量能、流动性与量价背离 =================
    "volume_ratio_5_20",
    "turnover_rate_mean_5",
    "turnover_deviation",
    "amihud_illiq_20",
    "turnover_cv_20",
    "pv_corr_20",
    "close_vwap_deviation",
    # ================= 5. 基本面财务特征 =================
    "roe",
    "roa",
    "profit_margin",
    "debt_to_equity",
    "current_ratio",
    "net_profit_yoy",
    "revenue_yoy",
    "healthy_expansion_velocity",
    # ================= 6. 基本面估值与截面动量共振 =================
    "EP",
    "BP",
    "CP",
    "market_cap_rank",
    "turnover_rank",
    "return_5_rank",
    "EP_rank",
    "pe_expansion_trend",
    "value_price_divergence",
    "active_market_cap",
    "ebit_rank",
]

# 因子定义（来自 regression.py）
FACTOR_DEFINITIONS = {}

# Label 定义
LABEL_FACTOR = {
    LABEL_NAME: "(ts_delay(close, -5) / ts_delay(open, -1)) - 1",
}

# 日期范围（探针筛选只在训练集上进行）
TRAIN_START = "20200101"
TRAIN_END = "20231231"
VAL_START = "20240101"
VAL_END = "20241231"


# 股票池筛选函数
def stock_pool_filter(df: pl.DataFrame) -> pl.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(1000, len(valid_df))
    small_cap_codes = valid_df.sort("total_mv").head(n)["ts_code"]
    return df["ts_code"].is_in(small_cap_codes)


def register_factors(
    engine: FactorEngine,
    selected_factors: List[str],
    factor_definitions: dict,
    label_factor: dict,
) -> List[str]:
    """注册因子"""
    print("=" * 80)
    print("注册因子")
    print("=" * 80)

    # 注册 SELECTED_FACTORS 中的因子（已在 metadata 中）
    print("\n注册特征因子（从 metadata）:")
    for name in selected_factors:
        engine.add_factor(name)
        print(f"  - {name}")

    # 注册 FACTOR_DEFINITIONS 中的因子（通过表达式）
    print("\n注册特征因子（表达式）:")
    for name, expr in factor_definitions.items():
        engine.add_factor(name, expr)
        print(f"  - {name}: {expr}")

    # 注册 label 因子
    print("\n注册 Label 因子（表达式）:")
    for name, expr in label_factor.items():
        engine.add_factor(name, expr)
        print(f"  - {name}: {expr}")

    feature_cols = selected_factors + list(factor_definitions.keys())

    print(f"\n特征因子数: {len(feature_cols)}")
    print(f"  - 来自 metadata: {len(selected_factors)}")
    print(f"  - 来自表达式: {len(factor_definitions)}")
    print(f"Label: {list(label_factor.keys())[0]}")
    print(f"已注册因子总数: {len(engine.list_registered())}")

    return feature_cols


def prepare_data_for_probe(
    engine: FactorEngine,
    feature_cols: List[str],
    start_date: str,
    end_date: str,
) -> pl.DataFrame:
    """准备探针筛选所需数据"""
    print("\n" + "=" * 80)
    print("准备探针筛选数据")
    print("=" * 80)

    factor_names = feature_cols + [LABEL_NAME]

    print(f"\n计算因子: {start_date} - {end_date}")
    data = engine.compute(
        factor_names=factor_names,
        start_date=start_date,
        end_date=end_date,
    )

    print(f"数据形状: {data.shape}")
    print(f"数据列: {len(data.columns)} 列")
    print(f"日期范围: {data['trade_date'].min()} - {data['trade_date'].max()}")

    return data


def apply_preprocessing_for_probe(
    data: pl.DataFrame,
    feature_cols: List[str],
) -> pl.DataFrame:
    """为探针筛选应用基础预处理（只处理缺失值）"""
    print("\n" + "=" * 80)
    print("数据预处理")
    print("=" * 80)

    # 只进行缺失值填充（避免标准化影响噪音分布）
    filler = NullFiller(feature_cols=feature_cols, strategy="mean")
    data = filler.fit_transform(data)

    print(f"缺失值处理完成")
    print(f"数据形状: {data.shape}")

    return data


def run_probe_feature_selection():
    """执行探针法因子筛选"""
    print("\n" + "=" * 80)
    print("增强探针法因子筛选")
    print("=" * 80)

    # 1. 创建 FactorEngine
    print("\n[1] 创建 FactorEngine")
    engine = FactorEngine()

    # 2. 注册因子
    print("\n[2] 注册因子")
    feature_cols = register_factors(
        engine, SELECTED_FACTORS, FACTOR_DEFINITIONS, LABEL_FACTOR
    )

    # 3. 准备数据（训练集 + 验证集，用于探针筛选）
    print("\n[3] 准备数据（训练集+验证集）")
    data = prepare_data_for_probe(
        engine=engine,
        feature_cols=feature_cols,
        start_date=TRAIN_START,
        end_date=VAL_END,  # 包含验证集，增加样本量
    )

    # 4. 股票池筛选
    print("\n[4] 执行股票池筛选")
    pool_manager = StockPoolManager(
        filter_func=stock_pool_filter,
        required_columns=["total_mv"],
        data_router=engine.router,
    )
    data = pool_manager.filter_and_select_daily(data)
    print(f"筛选后数据规模: {data.shape}")

    # 5. 数据预处理（只填充缺失值，不缩放）
    print("\n[5] 数据预处理")
    data = apply_preprocessing_for_probe(data, feature_cols)

    # 6. 数据质量检查
    print("\n[6] 数据质量检查")
    # check_data_quality(data, feature_cols, raise_on_error=True)
    print("[成功] 数据质量检查通过")

    # 7. 执行探针筛选
    print("\n[7] 执行探针筛选")
    selector = ProbeSelector(
        n_iterations=10,  # 迭代轮数
        n_noise_features=10,  # 每轮探针数
        validation_ratio=0.15,  # 验证集比例
        random_state=42,
        regression_params={
            "objective": "regression",
            "metric": "mae",
            "n_estimators": 200,
            "learning_rate": 0.05,
            "early_stopping_round": 30,
            "verbose": -1,
        },
        classification_params={
            "objective": "binary",
            "metric": "auc",
            "n_estimators": 200,
            "learning_rate": 0.05,
            "early_stopping_round": 30,
            "verbose": -1,
        },
        verbose=True,
    )

    selected_features = selector.select(
        data=data,
        feature_cols=feature_cols,
        target_col_regression=LABEL_NAME,
        date_col="trade_date",
    )

    # 8. 输出结果
    print("\n" + "=" * 80)
    print("探针筛选完成")
    print("=" * 80)
    print(f"\n原始特征数: {len(feature_cols)}")
    print(f"筛选后特征数: {len(selected_features)}")
    print(f"淘汰特征数: {len(feature_cols) - len(selected_features)}")
    print(
        f"淘汰比例: {(len(feature_cols) - len(selected_features)) / len(feature_cols):.1%}"
    )

    # 9. 保存筛选结果
    output_dir = "src/experiment/probe_selection/output"
    os.makedirs(output_dir, exist_ok=True)

    # 保存筛选后的特征列表
    output_file = os.path.join(output_dir, "selected_features.txt")
    with open(output_file, "w") as f:
        f.write("# 探针法筛选后的特征列表\n")
        f.write(f"# 原始特征数: {len(feature_cols)}\n")
        f.write(f"# 筛选后特征数: {len(selected_features)}\n")
        f.write(f"# 淘汰特征数: {len(feature_cols) - len(selected_features)}\n")
        f.write("\nSELECTED_FEATURES = [\n")
        for feat in selected_features:
            f.write(f'    "{feat}",\n')
        f.write("]\n")

    print(f"\n[保存] 筛选结果已保存到: {output_file}")

    # 10. 保存淘汰的特征
    eliminated_features = list(set(feature_cols) - set(selected_features))
    eliminated_file = os.path.join(output_dir, "eliminated_features.txt")
    with open(eliminated_file, "w") as f:
        f.write("# 被探针法淘汰的特征列表\n")
        f.write(f"# 淘汰总数: {len(eliminated_features)}\n")
        f.write("\nELIMINATED_FEATURES = [\n")
        for feat in eliminated_features:
            f.write(f'    "{feat}",\n')
        f.write("]\n")

    print(f"[保存] 淘汰特征已保存到: {eliminated_file}")

    # 11. 打印最终特征列表
    print("\n" + "=" * 80)
    print("最终特征列表（可直接复制到 regression.py）")
    print("=" * 80)
    print("\nSELECTED_FACTORS = [")
    for i, feat in enumerate(selected_features, 1):
        print(f'    "{feat}",')
    print("]")

    return selected_features


if __name__ == "__main__":
    selected = run_probe_feature_selection()
