# %% md
# ## 1. 导入依赖
# %%
import os
from datetime import datetime

import polars as pl

from src.factors import FactorEngine
from src.training import (
    DateSplitter,
    LightGBMModel,
    STFilter,
    StandardScaler,
    StockPoolManager,
    Trainer,
    Winsorizer,
    NullFiller,
    check_data_quality,
)
from src.training.config import TrainingConfig

# 从 common 模块导入共用配置和函数
from src.experiment.common import (
    SELECTED_FACTORS,
    FACTOR_DEFINITIONS,
    get_label_factor,
    register_factors,
    prepare_data,
    TRAIN_START,
    TRAIN_END,
    VAL_START,
    VAL_END,
    TEST_START,
    TEST_END,
    stock_pool_filter,
    STOCK_FILTER_REQUIRED_COLUMNS,
    OUTPUT_DIR,
    SAVE_PREDICTIONS,
    PERSIST_MODEL,
    TOP_N,
)


# %% md
# ## 2. 配置参数
#
# ### 2.1 标签定义
# %%
# Label 名称（回归任务使用连续收益率）
LABEL_NAME = "future_return_5"

# 获取 Label 因子定义
LABEL_FACTOR = get_label_factor(LABEL_NAME)

# 模型参数配置
MODEL_PARAMS = {
    "objective": "regression",
    "metric": "mae",  # 改为 MAE，对异常值更稳健
    # 树结构控制（防过拟合核心）
    # "num_leaves": 20,  # 从31降为20，降低模型复杂度
    # "max_depth": 16,  # 显式限制深度，防止过度拟合噪声
    # "min_child_samples": 50,  # 叶子最小样本数，防止学习极端样本
    # "min_child_weight": 0.001,
    # 学习参数
    "learning_rate": 0.01,  # 降低学习率，配合更多树
    "n_estimators": 1000,  # 增加树数量，配合早停
    # 采样策略（关键防过拟合）
    "subsample": 0.8,  # 每棵树随机采样80%数据（行采样）
    "subsample_freq": 5,  # 每5轮迭代进行一次 subsample
    "colsample_bytree": 0.8,  # 每棵树随机选择80%特征（列采样）
    # 正则化
    "reg_alpha": 0.1,  # L1正则，增加稀疏性
    "reg_lambda": 1.0,  # L2正则，平滑权重
    # 数值稳定性
    "verbose": -1,
    "random_state": 42,
}
# %% md
# ## 4. 训练流程
#
# ### 4.1 初始化组件
# %%
print("\n" + "=" * 80)
print("LightGBM 回归模型训练")
print("=" * 80)

# 1. 创建 FactorEngine（启用 metadata 功能）
print("\n[1] 创建 FactorEngine")
engine = FactorEngine(metadata_path="data/factors.jsonl")

# 2. 使用 metadata 定义因子
print("\n[2] 定义因子（从 metadata 注册）")
feature_cols = register_factors(
    engine, SELECTED_FACTORS, FACTOR_DEFINITIONS, LABEL_FACTOR
)
target_col = LABEL_NAME

# 3. 准备数据（使用模块级别的日期配置）
print("\n[3] 准备数据")

data = prepare_data(
    engine=engine,
    feature_cols=feature_cols,
    start_date=TRAIN_START,
    end_date=TEST_END,
    label_name=LABEL_NAME,
)

# 4. 打印配置信息
print(f"\n[配置] 训练期: {TRAIN_START} - {TRAIN_END}")
print(f"[配置] 验证期: {VAL_START} - {VAL_END}")
print(f"[配置] 测试期: {TEST_START} - {TEST_END}")
print(f"[配置] 特征数: {len(feature_cols)}")
print(f"[配置] 目标变量: {target_col}")

# 5. 创建模型
model = LightGBMModel(params=MODEL_PARAMS)

# 6. 创建数据处理器（使用函数返回的完整特征列表）
processors = [
    NullFiller(feature_cols=feature_cols, strategy="mean"),
    Winsorizer(feature_cols=feature_cols, lower=0.01, upper=0.99),
    StandardScaler(feature_cols=feature_cols),
]

# 7. 创建数据划分器（正确的 train/val/test 三分法）
# Train: 训练模型参数 | Val: 验证/早停 | Test: 最终评估
splitter = DateSplitter(
    train_start=TRAIN_START,
    train_end=TRAIN_END,
    val_start=VAL_START,
    val_end=VAL_END,
    test_start=TEST_START,
    test_end=TEST_END,
)

# 8. 创建股票池管理器
# 使用新的 API：传入自定义筛选函数和所需列
pool_manager = StockPoolManager(
    filter_func=stock_pool_filter,
    required_columns=STOCK_FILTER_REQUIRED_COLUMNS,  # 筛选所需的额外列
    # required_factors=STOCK_FILTER_REQUIRED_FACTORS,  # 可选：筛选所需的因子
    data_router=engine.router,
)
print("[股票池筛选] 使用自定义函数进行股票池筛选")
print(f"[股票池筛选] 所需基础列: {STOCK_FILTER_REQUIRED_COLUMNS}")
print("[股票池筛选] 筛选逻辑: 排除创业板/科创板/北交所后，每日选市值最小的500只")
# print(f"[股票池筛选] 所需因子: {list(STOCK_FILTER_REQUIRED_FACTORS.keys())}")

# 9. 创建 ST 股票过滤器
st_filter = STFilter(
    data_router=engine.router,
)

# 10. 创建训练器
trainer = Trainer(
    model=model,
    pool_manager=pool_manager,
    processors=processors,
    filters=[st_filter],  # 使用STFilter过滤ST股票
    splitter=splitter,
    target_col=target_col,
    feature_cols=feature_cols,
    persist_model=PERSIST_MODEL,
)
# %% md
# ### 4.2 执行训练
# %%
print("\n" + "=" * 80)
print("开始训练")
print("=" * 80)

# 步骤 1: 股票池筛选
print("\n[步骤 1/6] 股票池筛选")
print("-" * 60)
if pool_manager:
    print("  执行每日独立筛选股票池...")
    filtered_data = pool_manager.filter_and_select_daily(data)
    print(f"  筛选前数据规模: {data.shape}")
    print(f"  筛选后数据规模: {filtered_data.shape}")
    print(f"  筛选前股票数: {data['ts_code'].n_unique()}")
    print(f"  筛选后股票数: {filtered_data['ts_code'].n_unique()}")
    print(f"  删除记录数: {len(data) - len(filtered_data)}")
else:
    filtered_data = data
    print("  未配置股票池管理器，跳过筛选")
# %%
# 步骤 2: 划分训练/验证/测试集（正确的三分法）
print("\n[步骤 2/6] 划分训练集、验证集和测试集")
print("-" * 60)
if splitter:
    # 正确的三分法：train用于训练，val用于验证/早停，test仅用于最终评估
    train_data, val_data, test_data = splitter.split(filtered_data)
    print(f"  训练集数据规模: {train_data.shape}")
    print(f"  验证集数据规模: {val_data.shape}")
    print(f"  测试集数据规模: {test_data.shape}")
    print(f"  训练集股票数: {train_data['ts_code'].n_unique()}")
    print(f"  验证集股票数: {val_data['ts_code'].n_unique()}")
    print(f"  测试集股票数: {test_data['ts_code'].n_unique()}")
    print(
        f"  训练集日期范围: {train_data['trade_date'].min()} - {train_data['trade_date'].max()}"
    )
    print(
        f"  验证集日期范围: {val_data['trade_date'].min()} - {val_data['trade_date'].max()}"
    )
    print(
        f"  测试集日期范围: {test_data['trade_date'].min()} - {test_data['trade_date'].max()}"
    )

    print("\n  训练集前5行预览:")
    print(train_data.head())
    print("\n  验证集前5行预览:")
    print(val_data.head())
    print("\n  测试集前5行预览:")
    print(test_data.head())
else:
    train_data = filtered_data
    test_data = filtered_data
    print("  未配置划分器，全部作为训练集")
# %%
# 步骤 3: 数据质量检查（必须在预处理之前）
print("\n[步骤 3/7] 数据质量检查")
print("-" * 60)
print("  [说明] 此检查在 fillna 等处理之前执行，用于发现数据问题")

print("\n  检查训练集...")
check_data_quality(train_data, feature_cols, raise_on_error=True)

if "val_data" in locals() and val_data is not None:
    print("\n  检查验证集...")
    check_data_quality(val_data, feature_cols, raise_on_error=True)

print("\n  检查测试集...")
check_data_quality(test_data, feature_cols, raise_on_error=True)

print("  [成功] 数据质量检查通过，未发现异常")

# %%
# 步骤 4: 训练集数据处理
print("\n[步骤 4/7] 训练集数据处理")
print("-" * 60)
fitted_processors = []
if processors:
    for i, processor in enumerate(processors, 1):
        print(f"  [{i}/{len(processors)}] 应用处理器: {processor.__class__.__name__}")
        train_data_before = len(train_data)
        train_data = processor.fit_transform(train_data)
        train_data_after = len(train_data)
        fitted_processors.append(processor)
        print(f"      处理前记录数: {train_data_before}")
        print(f"      处理后记录数: {train_data_after}")
        if train_data_before != train_data_after:
            print(f"      删除记录数: {train_data_before - train_data_after}")

print("\n  训练集处理后前5行预览:")
print(train_data.head())
print(f"\n  训练集特征统计:")
print(f"    特征数: {len(feature_cols)}")
print(f"    样本数: {len(train_data)}")
print(f"    缺失值统计:")
for col in feature_cols[:5]:  # 只显示前5个特征的缺失值
    null_count = train_data[col].null_count()
    if null_count > 0:
        print(f"      {col}: {null_count} ({null_count / len(train_data) * 100:.2f}%)")
# %%
# 步骤 4: 训练模型
print("\n[步骤 5/7] 训练模型")
print("-" * 60)
print(f"  模型类型: LightGBM")
print(f"  训练样本数: {len(train_data)}")
print(f"  特征数: {len(feature_cols)}")
print(f"  目标变量: {target_col}")

X_train = train_data.select(feature_cols)
y_train = train_data.select(target_col).to_series()

print(f"\n  目标变量统计:")
print(f"    均值: {y_train.mean():.6f}")
print(f"    标准差: {y_train.std():.6f}")
print(f"    最小值: {y_train.min():.6f}")
print(f"    最大值: {y_train.max():.6f}")
print(f"    缺失值: {y_train.null_count()}")

print("\n  开始训练...")
model.fit(X_train, y_train)
print("  训练完成！")
# %%
# 步骤 5: 测试集数据处理
print("\n[步骤 6/7] 测试集数据处理")
print("-" * 60)
if processors and test_data is not train_data:
    for i, processor in enumerate(fitted_processors, 1):
        print(
            f"  [{i}/{len(fitted_processors)}] 应用处理器: {processor.__class__.__name__}"
        )
        test_data_before = len(test_data)
        test_data = processor.transform(test_data)
        test_data_after = len(test_data)
        print(f"      处理前记录数: {test_data_before}")
        print(f"      处理后记录数: {test_data_after}")
else:
    print("  跳过测试集处理")
# %%
# 步骤 6: 生成预测
print("\n[步骤 7/7] 生成预测")
print("-" * 60)
X_test = test_data.select(feature_cols)
print(f"  测试样本数: {len(X_test)}")
print("  预测中...")
predictions = model.predict(X_test)
print(f"  预测完成！")

print(f"\n  预测结果统计:")
print(f"    均值: {predictions.mean():.6f}")
print(f"    标准差: {predictions.std():.6f}")
print(f"    最小值: {predictions.min():.6f}")
print(f"    最大值: {predictions.max():.6f}")

# 保存结果到 trainer
trainer.results = test_data.with_columns([pl.Series("prediction", predictions)])
# %% md
# ### 4.3 训练指标曲线
# %%
print("\n" + "=" * 80)
print("训练指标曲线")
print("=" * 80)

# 重新训练以收集指标（因为之前的训练没有保存评估结果）
print("\n重新训练模型以收集训练指标...")

import lightgbm as lgb

# 准备数据（使用 val 做验证，test 不参与训练过程）
X_train_np = X_train.to_numpy()
y_train_np = y_train.to_numpy()
X_val_np = val_data.select(feature_cols).to_numpy()
y_val_np = val_data.select(target_col).to_series().to_numpy()

# 创建数据集
train_dataset = lgb.Dataset(X_train_np, label=y_train_np)
val_dataset = lgb.Dataset(X_val_np, label=y_val_np, reference=train_dataset)

# 用于存储评估结果
evals_result = {}

# 使用与原模型相同的参数重新训练
# 正确的三分法：train用于训练，val用于验证，test不参与训练过程
# 添加早停：如果验证指标连续100轮没有改善则停止训练
booster_with_eval = lgb.train(
    MODEL_PARAMS,
    train_dataset,
    num_boost_round=MODEL_PARAMS.get("n_estimators", 100),
    valid_sets=[train_dataset, val_dataset],
    valid_names=["train", "val"],
    callbacks=[
        lgb.record_evaluation(evals_result),
        lgb.early_stopping(stopping_rounds=100, verbose=True),
    ],
)

print("训练完成，指标已收集")

# 获取指标名称
metric_name = list(evals_result["train"].keys())[0]
print(f"\n评估指标: {metric_name}")

# 提取训练和验证指标
train_metric = evals_result["train"][metric_name]
val_metric = evals_result["val"][metric_name]

# 显示早停信息
actual_rounds = len(train_metric)
expected_rounds = MODEL_PARAMS.get("n_estimators", 100)
print(f"\n[早停信息]")
print(f"  配置的最大轮数: {expected_rounds}")
print(f"  实际训练轮数: {actual_rounds}")
if actual_rounds < expected_rounds:
    print(f"  早停状态: 已触发（连续100轮验证指标未改善）")
else:
    print(f"  早停状态: 未触发（达到最大轮数）")

print(f"\n最终指标:")
print(f"  训练 {metric_name}: {train_metric[-1]:.6f}")
print(f"  验证 {metric_name}: {val_metric[-1]:.6f}")
# %%
# 绘制训练指标曲线
import matplotlib.pyplot as plt

fig, ax = plt.subplots(figsize=(12, 6))

# 绘制训练集和验证集的指标曲线（注意：val用于验证，test不参与训练）
iterations = range(1, len(train_metric) + 1)
ax.plot(
    iterations, train_metric, label=f"Train {metric_name}", linewidth=2, color="blue"
)
ax.plot(
    iterations, val_metric, label=f"Validation {metric_name}", linewidth=2, color="red"
)

ax.set_xlabel("Iteration", fontsize=12)
ax.set_ylabel(metric_name.upper(), fontsize=12)
ax.set_title(
    f"Training and Validation {metric_name.upper()} Curve",
    fontsize=14,
    fontweight="bold",
)
ax.legend(fontsize=10)
ax.grid(True, alpha=0.3)

# 标记最佳验证指标点（用于早停决策）
best_iter = val_metric.index(min(val_metric))
best_metric = min(val_metric)
ax.axvline(
    x=best_iter + 1,
    color="green",
    linestyle="--",
    alpha=0.7,
    label=f"Best Iteration ({best_iter + 1})",
)
ax.scatter([best_iter + 1], [best_metric], color="green", s=100, zorder=5)
ax.annotate(
    f"Best: {best_metric:.6f}\nIter: {best_iter + 1}",
    xy=(best_iter + 1, best_metric),
    xytext=(best_iter + 1 + len(iterations) * 0.1, best_metric),
    fontsize=9,
    arrowprops=dict(arrowstyle="->", color="green", alpha=0.7),
)

plt.tight_layout()
plt.show()

print(f"\n[指标分析]")
print(f"  最佳验证 {metric_name}: {best_metric:.6f}")
print(f"  最佳迭代轮数: {best_iter + 1}")
print(f"  早停建议: 如果验证指标连续10轮不下降，建议在第 {best_iter + 1} 轮停止训练")
print(f"\n[重要提醒] 验证集仅用于早停/调参，测试集完全独立于训练过程！")
# %% md
# ### 4.4 查看结果
# %%
print("\n" + "=" * 80)
print("训练结果")
print("=" * 80)

results = trainer.results

print(f"\n结果数据形状: {results.shape}")
print(f"结果列: {results.columns}")
print(f"\n结果前10行预览:")
print(results.head(10))
print(f"\n结果后5行预览:")
print(results.tail())

print(f"\n每日预测样本数统计:")
daily_counts = results.group_by("trade_date").agg(pl.len()).sort("trade_date")
print(f"  最小: {daily_counts['len'].min()}")
print(f"  最大: {daily_counts['len'].max()}")
print(f"  平均: {daily_counts['len'].mean():.2f}")

# 展示某一天的前10个预测结果
sample_date = results["trade_date"][0]
sample_data = results.filter(results["trade_date"] == sample_date).head(10)
print(f"\n示例日期 {sample_date} 的前10条预测:")
print(sample_data.select(["ts_code", "trade_date", target_col, "prediction"]))
# %% md
# ### 4.4 保存结果
# %%
print("\n" + "=" * 80)
print("保存预测结果")
print("=" * 80)

# 确保输出目录存在
os.makedirs(OUTPUT_DIR, exist_ok=True)

# 生成时间戳
start_dt = datetime.strptime(TEST_START, "%Y%m%d")
end_dt = datetime.strptime(TEST_END, "%Y%m%d")
date_str = f"{start_dt.strftime('%Y%m%d')}_{end_dt.strftime('%Y%m%d')}"

# 保存每日 Top N
print(f"\n[1/1] 保存每日 Top {TOP_N} 股票...")
topn_output_path = os.path.join(OUTPUT_DIR, f"regression_output.csv")

# 按日期分组，取每日 top N
topn_by_date = []
unique_dates = results["trade_date"].unique().sort()
for date in unique_dates:
    day_data = results.filter(results["trade_date"] == date)
    # 按 prediction 降序排序，取前 N
    topn = day_data.sort("prediction", descending=True).head(TOP_N)
    topn_by_date.append(topn)

# 合并所有日期的 top N
topn_results = pl.concat(topn_by_date)

# 格式化日期并调整列顺序：日期、分数、股票
topn_to_save = topn_results.select(
    [
        pl.col("trade_date").str.slice(0, 4)
        + "-"
        + pl.col("trade_date").str.slice(4, 2)
        + "-"
        + pl.col("trade_date").str.slice(6, 2).alias("date"),
        pl.col("prediction").alias("score"),
        pl.col("ts_code"),
    ]
)
topn_to_save.write_csv(topn_output_path, include_header=True)
print(f"  保存路径: {topn_output_path}")
print(
    f"  保存行数: {len(topn_to_save)}（{len(unique_dates)}个交易日 × 每日top{TOP_N}）"
)
print(f"\n  预览（前15行）:")
print(topn_to_save.head(15))
# %% md
# ### 4.5 特征重要性
# %%
importance = model.feature_importance()
if importance is not None:
    print("\n特征重要性:")
    print(importance.sort_values(ascending=False))

print("\n" + "=" * 80)
print("训练完成！")
print("=" * 80)
# %% md
# ## 5. 可视化分析
#
# 使用训练好的模型直接绘图。
# - **特征重要性图**：辅助特征选择
# - **决策树图**：理解决策逻辑
# %%
# 导入可视化库
import matplotlib.pyplot as plt
import lightgbm as lgb
import pandas as pd

# 从封装的model中取出底层Booster
booster = model.model
print(f"模型类型: {type(booster)}")
print(f"特征数量: {len(feature_cols)}")
# %% md
# ### 5.1 绘制特征重要性（辅助特征选择）
#
# **解读**：
# - 重要性高的特征对模型贡献大
# - 重要性为0的特征可以考虑删除
# - 可以帮助理解哪些因子最有效
# %%
print("绘制特征重要性...")

fig, ax = plt.subplots(figsize=(10, 8))
lgb.plot_importance(
    booster,
    max_num_features=20,
    importance_type="gain",
    title="Feature Importance (Gain)",
    ax=ax,
)
ax.set_xlabel("Importance (Gain)")
plt.tight_layout()
plt.show()

# 打印重要性排名
importance_gain = pd.Series(
    booster.feature_importance(importance_type="gain"), index=feature_cols
).sort_values(ascending=False)

print("\n[特征重要性排名 - Gain]")
print(importance_gain)

# 识别低重要性特征
zero_importance = importance_gain[importance_gain == 0].index.tolist()
if zero_importance:
    print(f"\n[低重要性特征] 以下{len(zero_importance)}个特征重要性为0，可考虑删除:")
    for feat in zero_importance:
        print(f"  - {feat}")
else:
    print("\n所有特征都有一定重要性")