"""
Stage 2: LightGBM 抓非线性交互
滚动训练 + feature_fraction=0.8 + 分层收益验证
与 Ridge 对比：非线性是否提升了预测力？
"""
import pandas as pd
import numpy as np
from sklearn.preprocessing import StandardScaler
from sklearn.linear_model import Ridge
from lightgbm import LGBMRegressor
from scipy.stats import spearmanr

# ── 加载数据 ──
df = pd.read_parquet("data/btc_15m_binance.parquet")
df = df.set_index("ts").sort_index()

# ── 构造因子（只用 Stage 1 通过 IC>=0.02 的三个） ──
def make_factors(df):
    f = pd.DataFrame(index=df.index)
    f["ret_5"] = df["close"].pct_change(5)
    f["ret_20"] = df["close"].pct_change(20)
    f["ma_dev"] = df["close"] / df["close"].rolling(20).mean() - 1
    return f

# 标签
label = df["close"].shift(-4) / df["close"] - 1  # 未来1小时

# 合并
X_cols = ["ret_5","ret_20","ma_dev"]
data = pd.concat([make_factors(df), label.rename("fwd_ret")], axis=1).dropna()
print(f"样本: {len(data)}, 因子: {X_cols}")

# ── 滚动训练对比 ──
# 每月滚一次：训练≈30天，测试≈下1天
window_days = 30
step_days = 1
bars_per_day = 24 * 4

results = []

for start in range(0, len(data) - window_days * bars_per_day, step_days * bars_per_day):
    train_end = start + window_days * bars_per_day
    test_end = min(train_end + step_days * bars_per_day, len(data))
    
    train = data.iloc[start:train_end]
    test = data.iloc[train_end:test_end]
    if len(test) < 10:
        break
    
    X_train, y_train = train[X_cols], train["fwd_ret"]
    X_test, y_test = test[X_cols], test["fwd_ret"]
    
    scaler = StandardScaler()
    X_train_s = scaler.fit_transform(X_train)
    X_test_s = scaler.transform(X_test)
    
    # Ridge 基线
    ridge = Ridge(alpha=1.0).fit(X_train_s, y_train)
    ridge_pred = ridge.predict(X_test_s)
    ridge_ic, _ = spearmanr(ridge_pred, y_test)
    
    # LightGBM
    lgb = LGBMRegressor(
        n_estimators=100,
        max_depth=5,
        learning_rate=0.05,
        feature_fraction=0.8,   # 每棵树随机用80%特征 → 防过拟合
        subsample=0.8,
        random_state=42,
        verbosity=-1,
        force_col_wise=True,
    )
    lgb.fit(X_train_s, y_train)
    lgb_pred = lgb.predict(X_test_s)
    lgb_ic, _ = spearmanr(lgb_pred, y_test)
    
    results.append({
        "date": test.index[0],
        "ridge_ic": ridge_ic if not np.isnan(ridge_ic) else 0,
        "lgb_ic": lgb_ic if not np.isnan(lgb_ic) else 0,
    })

res = pd.DataFrame(results)
res["date"] = pd.to_datetime(res["date"])

# ── 结果 ──
print("\n" + "="*55)
print("Ridge(线性) vs LightGBM(非线性) IC 对比")
print("="*55)
print(f"Ridge    mean IC = {res['ridge_ic'].mean():+.4f}  std = {res['ridge_ic'].std():.4f}")
print(f"LightGBM mean IC = {res['lgb_ic'].mean():+.4f}  std = {res['lgb_ic'].std():.4f}")
print(f"\nRidge    IC>0: {(res['ridge_ic']>0).mean():.1%}")
print(f"LightGBM IC>0: {(res['lgb_ic']>0).mean():.1%}")

# ── 分层收益验证 ──
# 把预测值分5组，看最高组能否持续跑赢
def stratified_test(data, window_days=30, n_groups=5):
    """滚动分层：按预测值分5组，算每组平均收益"""
    bar_d = 24*4
    group_rets = {g: [] for g in range(n_groups)}
    
    for start in range(0, len(data) - window_days * bar_d, bar_d):
        train_end = start + window_days * bar_d
        test_end = min(train_end + bar_d, len(data))
        train = data.iloc[start:train_end]
        test = data.iloc[train_end:test_end]
        if len(test) < 10:
            break
        
        X_tr = StandardScaler().fit_transform(train[X_cols])
        X_te = StandardScaler().fit_transform(test[X_cols])
        
        lgb = LGBMRegressor(n_estimators=100, max_depth=5, learning_rate=0.05,
                            feature_fraction=0.8, subsample=0.8,
                            random_state=42, verbosity=-1, force_col_wise=True)
        lgb.fit(X_tr, train["fwd_ret"])
        pred = lgb.predict(X_te)
        
        test_w_pred = test.copy()
        test_w_pred["pred"] = pred
        test_w_pred["group"] = pd.qcut(test_w_pred["pred"], n_groups, 
                                        labels=False, duplicates="drop")
        
        for g in range(n_groups):
            grp = test_w_pred[test_w_pred["group"] == g]
            if len(grp) > 0:
                group_rets[g].append(grp["fwd_ret"].mean())
    
    return group_rets

print("\n" + "="*55)
print("分层收益验证（预测值分5组，看最高组能否跑赢）")
print("="*55)
group_rets = stratified_test(data)
for g in sorted(group_rets.keys()):
    rets = group_rets[g]
    mean_r = np.mean(rets) * 10000  # 转换为bps
    print(f"  Group {g} (最低→最高): avg ret = {mean_r:+.1f} bps, n={len(rets)}")

# 保存
res.to_parquet("data/stage2_compare.parquet")
print("\n📁 data/stage2_compare.parquet")