"""训练CatBoost二分类模型用于BTC 1小时涨跌预测"""

import sys
import os
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

import numpy as np
import pandas as pd
import pandas_ta as ta
from catboost import CatBoostClassifier, Pool
from sklearn.model_selection import train_test_split
from sklearn.metrics import classification_report, accuracy_score, confusion_matrix
from data.fetcher import BinanceDataFetcher
import pickle
from datetime import datetime


def add_technical_features(df):
    """添加丰富的技术指标特征"""
    df = df.copy()

    # 基础价格特征
    df['price_change'] = df['close'] - df['open']
    df['price_change_pct'] = (df['close'] - df['open']) / df['open']
    df['high_low_range'] = df['high'] - df['low']
    df['high_low_range_pct'] = (df['high'] - df['low']) / df['low']

    # 移动平均线
    df['sma_7'] = ta.sma(df['close'], length=7)
    df['sma_25'] = ta.sma(df['close'], length=25)
    df['sma_99'] = ta.sma(df['close'], length=99)
    df['ema_12'] = ta.ema(df['close'], length=12)
    df['ema_26'] = ta.ema(df['close'], length=26)

    # RSI
    df['rsi_14'] = ta.rsi(df['close'], length=14)
    df['rsi_7'] = ta.rsi(df['close'], length=7)

    # MACD
    macd = ta.macd(df['close'])
    if macd is not None and not macd.empty:
        df['macd'] = macd.iloc[:, 0]
        df['macd_hist'] = macd.iloc[:, 1]
        df['macd_signal'] = macd.iloc[:, 2]

    # 布林带
    bbands = ta.bbands(df['close'], length=20)
    if bbands is not None and not bbands.empty:
        df['bb_upper'] = bbands.iloc[:, 0]
        df['bb_mid'] = bbands.iloc[:, 1]
        df['bb_lower'] = bbands.iloc[:, 2]
        df['bb_width'] = (df['bb_upper'] - df['bb_lower']) / df['bb_mid']
        df['bb_position'] = (df['close'] - df['bb_lower']) / (df['bb_upper'] - df['bb_lower'])

    # ATR (Average True Range)
    df['atr'] = ta.atr(df['high'], df['low'], df['close'], length=14)

    # 成交量特征
    df['volume_sma_7'] = ta.sma(df['volume'], length=7)
    df['volume_ratio'] = df['volume'] / df['volume_sma_7']

    # OBV (On-Balance Volume)
    df['obv'] = ta.obv(df['close'], df['volume'])

    # Stochastic Oscillator
    stoch = ta.stoch(df['high'], df['low'], df['close'])
    if stoch is not None and not stoch.empty:
        df['stoch_k'] = stoch.iloc[:, 0]
        df['stoch_d'] = stoch.iloc[:, 1]

    # ADX (Average Directional Index)
    adx = ta.adx(df['high'], df['low'], df['close'])
    if adx is not None and not adx.empty:
        df['adx'] = adx.iloc[:, 0]

    # 价格动量
    df['momentum_1'] = df['close'].pct_change(1)
    df['momentum_3'] = df['close'].pct_change(3)
    df['momentum_7'] = df['close'].pct_change(7)

    return df


def create_labels(df, forecast_horizon=1):
    """创建二分类标签"""
    df = df.copy()
    df['future_price'] = df['close'].shift(-forecast_horizon)
    df['future_return'] = (df['future_price'] - df['close']) / df['close']
    df['label'] = (df['future_return'] > 0).astype(int)  # 1=涨, 0=跌
    return df[:-forecast_horizon]


def prepare_data(df):
    """准备训练数据"""
    # 添加技术指标
    df = add_technical_features(df)

    # 创建标签
    df = create_labels(df, forecast_horizon=1)

    # 删除NaN
    df = df.dropna().reset_index(drop=True)

    # 选择特征列（排除原始时间戳和标签列）
    exclude_cols = ['timestamp', 'label', 'future_price', 'future_return']
    feature_cols = [col for col in df.columns if col not in exclude_cols]

    X = df[feature_cols]
    y = df['label']

    return X, y, feature_cols


def main():
    print("="*80)
    print("训练CatBoost二分类模型 - BTC 1小时涨跌预测")
    print("="*80)

    # ========== 步骤1: 获取数据 ==========
    print("\n[1/5] 获取历史数据...")

    fetcher = BinanceDataFetcher()

    # 获取训练数据 (2020-2024)
    df_train = fetcher.fetch_ohlcv('BTC', '1h', start_date='2020-01-01',
                                    end_date='2024-12-31', use_cache=True)

    # 获取验证数据 (2025)
    df_val = fetcher.fetch_ohlcv('BTC', '1h', start_date='2025-01-01',
                                  end_date='2025-12-31', use_cache=True)

    print(f"✓ 训练数据: {len(df_train)} 行 ({df_train['timestamp'].min()} 至 {df_train['timestamp'].max()})")
    print(f"✓ 验证数据: {len(df_val)} 行 ({df_val['timestamp'].min()} 至 {df_val['timestamp'].max()})")

    # ========== 步骤2: 特征工程 ==========
    print("\n[2/5] 特征工程...")

    X_train, y_train, feature_cols = prepare_data(df_train)
    X_val, y_val, _ = prepare_data(df_val)

    print(f"✓ 特征数量: {len(feature_cols)}")
    print(f"✓ 训练样本: {len(X_train)} (涨: {(y_train==1).sum()}, 跌: {(y_train==0).sum()})")
    print(f"✓ 验证样本: {len(X_val)} (涨: {(y_val==1).sum()}, 跌: {(y_val==0).sum()})")

    # 显示部分特征
    print(f"\n特征列表 (前20个):")
    for i, col in enumerate(feature_cols[:20], 1):
        print(f"  {i:2d}. {col}")
    if len(feature_cols) > 20:
        print(f"  ... 还有 {len(feature_cols)-20} 个特征")

    # ========== 步骤3: 训练CatBoost模型 ==========
    print("\n[3/5] 训练CatBoost模型...")

    # 创建CatBoost数据池
    train_pool = Pool(X_train, y_train)
    val_pool = Pool(X_val, y_val)

    # 配置CatBoost参数
    model = CatBoostClassifier(
        iterations=1000,
        learning_rate=0.05,
        depth=6,
        l2_leaf_reg=3,
        loss_function='Logloss',
        eval_metric='Accuracy',
        random_seed=42,
        verbose=100,
        early_stopping_rounds=50,
        task_type='CPU',  # 如果有GPU可改为'GPU'
    )

    # 训练模型
    model.fit(
        train_pool,
        eval_set=val_pool,
        use_best_model=True,
        plot=False
    )

    print(f"\n✓ 训练完成")
    print(f"  最佳迭代轮次: {model.get_best_iteration()}")
    print(f"  最佳验证准确率: {model.get_best_score()['validation']['Accuracy']:.4f}")

    # ========== 步骤4: 评估模型 ==========
    print("\n[4/5] 评估模型性能...")

    # 训练集预测
    y_train_pred = model.predict(X_train)
    train_acc = accuracy_score(y_train, y_train_pred)

    # 验证集预测
    y_val_pred = model.predict(X_val)
    val_acc = accuracy_score(y_val, y_val_pred)

    print(f"\n训练集性能：")
    print(f"  准确率: {train_acc*100:.2f}%")

    print(f"\n验证集性能：")
    print(f"  准确率: {val_acc*100:.2f}%")

    # 混淆矩阵
    cm = confusion_matrix(y_val, y_val_pred)
    print(f"\n混淆矩阵 (验证集):")
    print(f"  预测↓  预测↑")
    print(f"真实↓ {cm[0,0]:4d}   {cm[0,1]:4d}")
    print(f"真实↑ {cm[1,0]:4d}   {cm[1,1]:4d}")

    # 分类报告
    print(f"\n分类报告 (验证集):")
    print(classification_report(y_val, y_val_pred, target_names=['下跌', '上涨']))

    # 特征重要性 (Top 20)
    feature_importance = model.get_feature_importance()
    importance_df = pd.DataFrame({
        'feature': feature_cols,
        'importance': feature_importance
    }).sort_values('importance', ascending=False)

    print(f"\nTop 20 重要特征:")
    for i, (idx, row) in enumerate(importance_df.head(20).iterrows(), 1):
        print(f"  {i:2d}. {row['feature']:<20s} {row['importance']:7.2f}")

    # ========== 步骤5: 保存模型 ==========
    print("\n[5/5] 保存模型...")

    checkpoint_dir = './models/checkpoints/catboost_classification_2026'
    os.makedirs(checkpoint_dir, exist_ok=True)

    # 保存CatBoost模型
    model_path = f'{checkpoint_dir}/catboost_model.cbm'
    model.save_model(model_path, format='cbm')

    # 保存特征列表和元数据
    metadata = {
        'feature_cols': feature_cols,
        'train_acc': float(train_acc),
        'val_acc': float(val_acc),
        'best_iteration': int(model.get_best_iteration()),
        'train_samples': len(X_train),
        'val_samples': len(X_val),
        'timestamp': datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
        'train_date_range': f"{df_train['timestamp'].min()} - {df_train['timestamp'].max()}",
        'val_date_range': f"{df_val['timestamp'].min()} - {df_val['timestamp'].max()}",
    }

    with open(f'{checkpoint_dir}/metadata.pkl', 'wb') as f:
        pickle.dump(metadata, f)

    print(f"✓ 模型已保存到: {model_path}")
    print(f"✓ 元数据已保存到: {checkpoint_dir}/metadata.pkl")

    # ========== 总结 ==========
    print("\n" + "="*80)
    print("✓ CatBoost二分类模型训练完成!")
    print("="*80)

    print(f"\n模型性能总结：")
    print(f"  训练集准确率: {train_acc*100:.2f}%")
    print(f"  验证集准确率: {val_acc*100:.2f}%")
    print(f"  特征数量: {len(feature_cols)}")
    print(f"  最佳迭代: {model.get_best_iteration()}")

    print(f"\n保存位置：")
    print(f"  {checkpoint_dir}/")

    print("\n" + "="*80)


if __name__ == "__main__":
    main()