# %% md
# # TabPFN 回归训练流程
#
# 使用 TabPFN (Prior-Data Fitted Network) 进行回归预测。
# TabPFN 通过上下文学习进行预测，无需传统梯度下降训练过程。
# %% md
# ## 1. 导入依赖
# %%
import os

from src.factors import FactorEngine
from src.training import (
    FactorManager,
    DataPipeline,
    NullFiller,
    Winsorizer,
    StandardScaler,
    CrossSectionalStandardScaler,
)
from src.training.core.trainer_v2 import Trainer
from src.training.components.filters import STFilter
from src.training.components.models import TabPFNModel
from src.experiment.common import (
    SELECTED_FACTORS,
    FACTOR_DEFINITIONS,
    LABEL_NAME,
    LABEL_FACTOR,
    TRAIN_START,
    TRAIN_END,
    VAL_START,
    VAL_END,
    TEST_START,
    TEST_END,
    stock_pool_filter,
    STOCK_FILTER_REQUIRED_COLUMNS,
    OUTPUT_DIR,
    SAVE_PREDICTIONS,
    SAVE_MODEL,
    get_model_save_path,
    save_model_with_factors,
    TOP_N,
)

# 训练类型标识
TRAINING_TYPE = "tabpfn"

# %% md
# ## 2. 训练特定配置
# %%
# Label 配置（从 common.py 统一导入）
# LABEL_NAME 和 LABEL_FACTOR 已在 common.py 中绑定，只需从 common 导入

# 排除的因子列表（与 regression.py 保持一致）
EXCLUDED_FACTORS = [
    "GTJA_alpha001",
    "GTJA_alpha002",
    "GTJA_alpha003",
    "GTJA_alpha004",
    "GTJA_alpha005",
    "GTJA_alpha006",
    "GTJA_alpha007",
    "GTJA_alpha008",
    "GTJA_alpha009",
    "GTJA_alpha010",
    "GTJA_alpha011",
    "GTJA_alpha012",
    "GTJA_alpha013",
    "GTJA_alpha014",
    "GTJA_alpha015",
    "GTJA_alpha016",
    "GTJA_alpha017",
    "GTJA_alpha018",
    "GTJA_alpha019",
    "GTJA_alpha020",
    "GTJA_alpha022",
    "GTJA_alpha023",
    "GTJA_alpha024",
    "GTJA_alpha025",
    "GTJA_alpha026",
    "GTJA_alpha027",
    "GTJA_alpha028",
    "GTJA_alpha029",
    "GTJA_alpha031",
    "GTJA_alpha032",
    "GTJA_alpha033",
    "GTJA_alpha034",
    "GTJA_alpha035",
    "GTJA_alpha036",
    "GTJA_alpha037",
    # "GTJA_alpha038",
    "GTJA_alpha039",
    "GTJA_alpha040",
    "GTJA_alpha041",
    "GTJA_alpha042",
    "GTJA_alpha043",
    "GTJA_alpha044",
    "GTJA_alpha045",
    "GTJA_alpha046",
    "GTJA_alpha047",
    "GTJA_alpha048",
    "GTJA_alpha049",
    "GTJA_alpha050",
    "GTJA_alpha051",
    "GTJA_alpha052",
    "GTJA_alpha053",
    "GTJA_alpha054",
    "GTJA_alpha056",
    "GTJA_alpha057",
    "GTJA_alpha058",
    "GTJA_alpha059",
    "GTJA_alpha060",
    "GTJA_alpha061",
    "GTJA_alpha062",
    "GTJA_alpha063",
    "GTJA_alpha064",
    "GTJA_alpha065",
    "GTJA_alpha066",
    "GTJA_alpha067",
    "GTJA_alpha068",
    "GTJA_alpha070",
    "GTJA_alpha071",
    "GTJA_alpha072",
    "GTJA_alpha073",
    "GTJA_alpha074",
    "GTJA_alpha076",
    "GTJA_alpha077",
    "GTJA_alpha078",
    "GTJA_alpha079",
    "GTJA_alpha080",
    "GTJA_alpha081",
    "GTJA_alpha082",
    "GTJA_alpha083",
    "GTJA_alpha084",
    "GTJA_alpha085",
    "GTJA_alpha086",
    "GTJA_alpha087",
    "GTJA_alpha088",
    "GTJA_alpha089",
    "GTJA_alpha090",
    "GTJA_alpha091",
    "GTJA_alpha092",
    "GTJA_alpha093",
    "GTJA_alpha094",
    "GTJA_alpha095",
    "GTJA_alpha096",
    "GTJA_alpha097",
    "GTJA_alpha098",
    "GTJA_alpha099",
    "GTJA_alpha100",
    "GTJA_alpha101",
    "GTJA_alpha102",
    "GTJA_alpha103",
    "GTJA_alpha104",
    "GTJA_alpha105",
    "GTJA_alpha106",
    "GTJA_alpha107",
    "GTJA_alpha108",
    "GTJA_alpha109",
    "GTJA_alpha110",
    "GTJA_alpha111",
    "GTJA_alpha112",
    # "GTJA_alpha113",
    "GTJA_alpha114",
    "GTJA_alpha115",
    "GTJA_alpha117",
    "GTJA_alpha118",
    "GTJA_alpha119",
    "GTJA_alpha120",
    # "GTJA_alpha121",
    "GTJA_alpha122",
    "GTJA_alpha123",
    "GTJA_alpha124",
    "GTJA_alpha125",
    "GTJA_alpha126",
    "GTJA_alpha127",
    "GTJA_alpha128",
    "GTJA_alpha129",
    "GTJA_alpha130",
    "GTJA_alpha131",
    "GTJA_alpha132",
    "GTJA_alpha133",
    "GTJA_alpha134",
    "GTJA_alpha135",
    "GTJA_alpha136",
    # "GTJA_alpha138",
    "GTJA_alpha139",
    # "GTJA_alpha140",
    "GTJA_alpha141",
    "GTJA_alpha142",
    "GTJA_alpha145",
    # "GTJA_alpha146",
    "GTJA_alpha148",
    "GTJA_alpha150",
    "GTJA_alpha151",
    "GTJA_alpha152",
    "GTJA_alpha153",
    "GTJA_alpha154",
    "GTJA_alpha155",
    "GTJA_alpha156",
    "GTJA_alpha157",
    "GTJA_alpha158",
    "GTJA_alpha159",
    "GTJA_alpha160",
    "GTJA_alpha161",
    "GTJA_alpha162",
    "GTJA_alpha163",
    "GTJA_alpha164",
    # "GTJA_alpha165",
    "GTJA_alpha166",
    "GTJA_alpha167",
    "GTJA_alpha168",
    "GTJA_alpha169",
    "GTJA_alpha170",
    "GTJA_alpha171",
    "GTJA_alpha173",
    "GTJA_alpha174",
    "GTJA_alpha175",
    "GTJA_alpha176",
    "GTJA_alpha177",
    "GTJA_alpha178",
    "GTJA_alpha179",
    "GTJA_alpha180",
    # "GTJA_alpha183",
    "GTJA_alpha184",
    "GTJA_alpha185",
    "GTJA_alpha187",
    "GTJA_alpha188",
    "GTJA_alpha189",
    "GTJA_alpha191",
    "chip_dispersion_90",
    "chip_dispersion_70",
    "cost_skewness",
    "dispersion_change_20",
    "price_to_avg_cost",
    "price_to_median_cost",
    "mean_median_dev",
    "trap_pressure",
    "bottom_profit",
    "history_position",
    "winner_rate_surge_5",
    "winner_rate_cs_rank",
    "winner_rate_dev_20",
    "winner_rate_volatility",
    "smart_money_accumulation",
    "winner_vol_corr_20",
    "cost_base_momentum",
    "bottom_cost_stability",
    "pivot_reversion",
    "chip_transition",
]

# 模型参数配置
MODEL_PARAMS = {
    # ==================== 设备配置 ====================
    "device": "cuda",  # 计算设备: "cuda" 或 "cpu"（默认 cuda）
    # ==================== 上下文限制 ====================
    "max_context_size": 100,  # 16GB GPU 建议 1000-3000，32GB 可尝试 5000-8000
}

# 日期范围配置
date_range = {
    "train": (TRAIN_START, TRAIN_END),
    "val": (VAL_START, VAL_END),
    "test": (TEST_START, TEST_END),
}

# 输出配置
output_config = {
    "output_dir": OUTPUT_DIR,
    "output_filename": "tabpfn_output.csv",
    "save_predictions": SAVE_PREDICTIONS,
    "save_model": SAVE_MODEL,
    "model_save_path": get_model_save_path(TRAINING_TYPE),
    "top_n": TOP_N,
}


# %% md
# ## 3. 自定义 TabPFN 任务
# %%
from src.training.tasks import RegressionTask


class TabPFNTask(RegressionTask):
    """TabPFN 回归任务

    继承自 RegressionTask，但使用 TabPFNModel 作为模型。
    TabPFN 不需要传统的训练过程，而是通过上下文学习进行预测。
    """

    def __init__(self, model_params: dict, label_name: str):
        """初始化 TabPFN 任务

        Args:
            model_params: TabPFN 参数字典
            label_name: Label 列名称
        """
        # 不调用父类 __init__，直接初始化以避免创建 LightGBMModel
        from src.training.tasks.base import BaseTask

        BaseTask.__init__(self, model_params, label_name)
        self.evals_result: dict | None = None
        self.model = TabPFNModel(params=model_params)

    def fit(self, train_data: dict, val_data: dict) -> None:
        """训练 TabPFN 模型

        TabPFN 通过将训练数据加载到模型上下文中进行"训练"，
        不需要传统的梯度下降优化过程。

        Args:
            train_data: 训练数据 {"X": DataFrame, "y": Series}
            val_data: 验证数据，用于评估但不参与训练
        """
        X_train = train_data["X"]
        y_train = train_data["y"]
        X_val = val_data.get("X")
        y_val = val_data.get("y")

        # TabPFN 使用 eval_set 进行验证
        self.model.fit(
            X_train, y_train, eval_set=(X_val, y_val) if X_val is not None else None
        )

    def get_model(self) -> TabPFNModel:
        """获取训练好的模型实例"""
        return self.model


# %% md
# ## 4. 主函数
# %%
def main():
    """主函数"""
    print("\n" + "=" * 80)
    print("TabPFN 回归模型训练")
    print("=" * 80)
    print("\n[说明] TabPFN 使用上下文学习(In-Context Learning)，")
    print("       训练过程实际是加载数据到模型上下文。")
    print("       如果训练数据超过上下文限制，会自动截取最近的数据。")

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

    # 2. 创建 FactorManager
    print("\n[2] 创建 FactorManager")
    factor_manager = FactorManager(
        selected_factors=SELECTED_FACTORS,
        factor_definitions=FACTOR_DEFINITIONS,
        label_factor=LABEL_FACTOR,
        excluded_factors=EXCLUDED_FACTORS,
    )

    # 3. 创建 DataPipeline
    print("\n[3] 创建 DataPipeline")
    pipeline = DataPipeline(
        factor_manager=factor_manager,
        processor_configs=[
            (NullFiller, {"strategy": "mean"}),
            (Winsorizer, {"lower": 0.01, "upper": 0.99}),
            (StandardScaler, {}),
            # (CrossSectionalStandardScaler, {}),
        ],
        label_processor_configs=[
            # 对 label 进行缩尾处理（去除极端收益率）
            (Winsorizer, {"lower": 0.05, "upper": 0.95}),
            # (StandardScaler, {}),
        ],
        filters=[STFilter(data_router=engine.router)],
        stock_pool_filter_func=stock_pool_filter,
        stock_pool_required_columns=STOCK_FILTER_REQUIRED_COLUMNS,
    )

    # 4. 创建 TabPFNTask
    print("\n[4] 创建 TabPFNTask")
    task = TabPFNTask(
        model_params=MODEL_PARAMS,
        label_name=LABEL_NAME,
    )

    # 5. 创建 Trainer
    print("\n[5] 创建 Trainer")
    trainer = Trainer(
        data_pipeline=pipeline,
        task=task,
        output_config=output_config,
        verbose=True,
    )

    # 6. 执行训练
    print("\n[6] 执行训练")
    results = trainer.run(engine=engine, date_range=date_range)

    # 7. 保存模型和因子信息（如果启用）
    if SAVE_MODEL:
        print("\n[7] 保存模型和因子信息")
        save_model_with_factors(
            model=task.get_model(),
            model_path=output_config["model_save_path"],
            selected_factors=SELECTED_FACTORS,
            factor_definitions=FACTOR_DEFINITIONS,
            fitted_processors=pipeline.get_fitted_processors(),
        )

    # 8. 输出 TabPFN 特有指标
    print("\n" + "=" * 80)
    print("TabPFN 训练完成！")
    print(f"结果保存路径: {os.path.join(OUTPUT_DIR, 'tabpfn_output.csv')}")

    # 显示验证集评估结果（如果可用）
    model = task.get_model()
    best_score = model.get_best_score()
    if best_score:
        print("\n[验证集评估指标]")
        for metric, value in best_score.get("valid_0", {}).items():
            print(f"  - {metric}: {value:.6f}")

    print("=" * 80)

    return results


if __name__ == "__main__":
    main()
