"""
Stage 1: Ridge 线性回归 + IC 验证
5个价量因子 → Ridge打分 → 计算信息系数(IC)
IC ≥ 0.02 视为有效因子
"""
import pandas as pd
import numpy as np
from sklearn.linear_model import Ridge
from sklearn.preprocessing import StandardScaler
from scipy.stats import spearmanr

# ── 加载数据 ──
df = pd.read_parquet("data/btc_15m_binance.parquet")
df = df.set_index("ts").sort_index()
print(f"数据: {len(df)} 根K线, {df.index.min()} → {df.index.max()}")

# ── 构造因子 ──
def make_factors(df):
    f = pd.DataFrame(index=df.index)
    
    # 因子1: 5根K线收益率
    f["ret_5"] = df["close"].pct_change(5)
    
    # 因子2: 20根K线收益率
    f["ret_20"] = df["close"].pct_change(20)
    
    # 因子3: 20根K线波动率
    f["vol_20"] = df["close"].pct_change().rolling(20).std()
    
    # 因子4: 均线乖离 (close/MA20 - 1)
    f["ma_dev"] = df["close"] / df["close"].rolling(20).mean() - 1
    
    # 因子5: 成交量变化 (volume/vol_ma20 - 1)
    f["vol_chg"] = df["volume"] / df["volume"].rolling(20).mean() - 1
    
    return f

# ── 标签: 未来N根K线收益率 ──
def make_label(df, horizon=4):
    """horizon=4 → 未来1小时收益 (4×15min)"""
    return df["close"].shift(-horizon) / df["close"] - 1

# ── 主流程 ──
factors = make_factors(df)
label = make_label(df, horizon=4)

# 合并并去NaN
data = pd.concat([factors, label.rename("fwd_ret")], axis=1).dropna()
print(f"有效样本: {len(data)}")

# ── 滚动IC计算 ──
# 用前30天训练，预测下一天，逐日滚动
window = 30 * 24 * 4  # 30天的15分钟K线数
step = 24 * 4          # 每天滚动一次

ics = []
dates = []
X_cols = ["ret_5","ret_20","vol_20","ma_dev","vol_chg"]

for i in range(0, len(data) - window - step, step):
    train = data.iloc[i:i+window]
    test = data.iloc[i+window:i+window+step]
    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回归
    model = Ridge(alpha=1.0)
    model.fit(X_train_s, y_train)
    
    # 预测
    pred = model.predict(X_test_s)
    
    # IC (Spearman rank correlation)
    ic, pval = spearmanr(pred, y_test)
    if not np.isnan(ic):
        ics.append(ic)
        dates.append(test.index[0])

ic_df = pd.DataFrame({"date": dates, "ic": ics})
ic_df["date"] = pd.to_datetime(ic_df["date"])

# ── 逐因子IC ──
print("\n" + "="*50)
print("逐因子IC (单因子+标准化 → Ridge → Spearman IC)")
print("="*50)

factor_ics = {}
for col in X_cols:
    col_ics = []
    for i in range(0, len(data) - window - step, step):
        train = data.iloc[i:i+window]
        test = data.iloc[i+window:i+window+step]
        if len(test) < 10:
            break
        X_tr = StandardScaler().fit_transform(train[[col]])
        X_te = StandardScaler().fit_transform(test[[col]])
        m = Ridge(alpha=1.0).fit(X_tr, train["fwd_ret"])
        pred = m.predict(X_te)
        ic, _ = spearmanr(pred, test["fwd_ret"])
        if not np.isnan(ic):
            col_ics.append(ic)
    mean_ic = np.mean(col_ics)
    factor_ics[col] = mean_ic
    status = "✅" if abs(mean_ic) >= 0.02 else "❌"
    print(f"  {col:12s}  mean IC = {mean_ic:+.4f}  {status}")

# ── 汇总 ──
print(f"\n整体 IC: mean={ic_df['ic'].mean():+.4f}, std={ic_df['ic'].std():.4f}")
print(f"IC>0 占比: {(ic_df['ic']>0).mean():.1%}")
print(f"IC>=0.02 占比: {(ic_df['ic'].abs()>=0.02).mean():.1%}")

# 保存
ic_df.to_parquet("data/stage1_ic.parquet")
pd.Series(factor_ics, name="mean_ic").to_csv("data/stage1_factor_ic.csv")
print("\n📁 结果已保存: data/stage1_ic.parquet, data/stage1_factor_ic.csv")