jyx_code4/strategy/ai_strategy.py

"""
方案B：AI模型训练 + 信号生成
使用 LightGBM / XGBoost，Walk-Forward 滚动训练
"""
import json
import joblib
import pandas as pd
import numpy as np
from pathlib import Path
from loguru import logger

from .config import MODEL_CONFIG as MC, PRIMARY_PERIOD, PROJECT_ROOT
from .feature_engine import prepare_dataset, get_latest_feature_row
from .backtest import BacktestEngine, print_metrics

SCHEME_B_MODEL_DIR = PROJECT_ROOT / 'models'
SCHEME_B_MODEL_FILE = SCHEME_B_MODEL_DIR / 'scheme_b_last_model.joblib'
SCHEME_B_SCALER_FILE = SCHEME_B_MODEL_DIR / 'scheme_b_scaler.joblib'
SCHEME_B_FEATURES_FILE = SCHEME_B_MODEL_DIR / 'scheme_b_features.json'


class AIStrategy:
    """AI模型策略 — LightGBM / XGBoost Walk-Forward"""

    def __init__(self, model_type: str = 'lightgbm'):
        """
        :param model_type: 'lightgbm' 或 'xgboost'
        """
        self.model_type = model_type
        self.models = []  # 存储每个窗口训练的模型
        self.feature_importance = None

    def _create_model(self):
        """创建模型实例"""
        if self.model_type == 'lightgbm':
            import lightgbm as lgb
            params = MC['lightgbm'].copy()
            return lgb.LGBMClassifier(**params)
        elif self.model_type == 'xgboost':
            import xgboost as xgb
            params = MC['xgboost'].copy()
            return xgb.XGBClassifier(**params)
        else:
            raise ValueError(f"不支持的模型类型: {self.model_type}")

    def walk_forward_train(self, X: pd.DataFrame, y: pd.Series,
                           confidence_threshold: float = 0.45) -> pd.Series:
        """
        Walk-Forward 滚动训练与预测
        :param confidence_threshold: 概率阈值，低于此值的预测设为0（观望）
        :return: 全部测试窗口拼接的预测信号
        """
        train_size = MC['walk_forward_train_size']
        test_size = MC['walk_forward_test_size']
        step = MC['walk_forward_step']

        n = len(X)
        all_preds = pd.Series(dtype=float)
        window_count = 0

        logger.info(f"Walk-Forward: 数据量={n}, 训练窗口={train_size}, "
                     f"测试窗口={test_size}, 步长={step}, 置信阈值={confidence_threshold}")

        start = 0
        while start + train_size + test_size <= n:
            train_end = start + train_size
            test_end = min(train_end + test_size, n)

            X_train = X.iloc[start:train_end]
            y_train = y.iloc[start:train_end]
            X_test = X.iloc[train_end:test_end]
            y_test = y.iloc[train_end:test_end]

            # 训练
            model = self._create_model()
            model.fit(X_train, y_train)
            self.models.append(model)

            # 预测概率 + 置信度过滤
            proba = model.predict_proba(X_test)
            max_proba = proba.max(axis=1)
            raw_preds = model.predict(X_test)

            # 置信度不够的设为观望
            filtered_preds = raw_preds.copy()
            filtered_preds[max_proba < confidence_threshold] = 0

            preds = pd.Series(filtered_preds, index=X_test.index)
            all_preds = pd.concat([all_preds, preds])

            # 准确率（用原始预测算）
            acc = (raw_preds == y_test).mean()
            n_filtered = (max_proba < confidence_threshold).sum()
            window_count += 1
            logger.info(f"  窗口 {window_count}: 训练[{start}:{train_end}] "
                         f"测试[{train_end}:{test_end}] 准确率={acc:.2%} "
                         f"过滤={n_filtered}/{len(X_test)}")

            start += step

        # 特征重要性（取最后一个模型）
        if self.models:
            last_model = self.models[-1]
            if hasattr(last_model, 'feature_importances_'):
                self.feature_importance = pd.Series(
                    last_model.feature_importances_, index=X.columns
                ).sort_values(ascending=False)

        logger.info(f"Walk-Forward 完成: {window_count} 个窗口, "
                     f"共 {len(all_preds)} 条预测")
        return all_preds

    def get_top_features(self, n: int = 20) -> pd.Series:
        """获取Top N重要特征"""
        if self.feature_importance is not None:
            return self.feature_importance.head(n)
        return pd.Series(dtype=float)

    def run(self, period: int = None, start_date: str = None, end_date: str = None) -> dict:
        """
        完整运行方案B
        若指定了 start_date/end_date，会向前加载 warm_up_months 月数据用于训练，使回测区间首月即有预测。
        :return: 回测结果
        """
        if period is None:
            period = PRIMARY_PERIOD

        logger.info("=" * 60)
        logger.info(f"方案B：AI模型策略 ({self.model_type})")
        logger.info("=" * 60)

        from .data_loader import load_kline

        # 1. 准备数据：若指定了回测区间，则向前加载预热数据，使区间内从首月就有预测
        load_start, load_end = start_date, end_date
        if start_date and end_date:
            warm_months = MC.get('warm_up_months', 12)
            load_start_ts = pd.Timestamp(start_date) - pd.DateOffset(months=warm_months)
            load_start = load_start_ts.strftime('%Y-%m-%d')
            logger.info(f"回测区间 [{start_date} ~ {end_date}]，向前加载 {warm_months} 月至 {load_start} 用于训练")

        X, y, feature_names, scaler = prepare_dataset(period, load_start, load_end)

        # 2. Walk-Forward 训练
        predictions = self.walk_forward_train(X, y)

        # 3. 回测仅用用户指定区间；将预测对齐到该区间的每根K线
        df = load_kline(period, start_date, end_date)
        if df.empty:
            logger.warning("回测区间内无K线数据")
            return BacktestEngine()._empty_result()

        # 对齐信号：回测区间内有的时间戳用预测，缺失的填 0（观望）
        signals = predictions.reindex(df.index, fill_value=0).astype(int)
        prices = df['close']

        # 4. 回测
        engine = BacktestEngine()
        result = engine.run(prices, signals)

        print_metrics(result['metrics'], f"方案B: {self.model_type} AI策略")

        # 5. 保存最后一窗模型、scaler、特征列（供实盘 get_live_signal 使用）
        if self.models and scaler is not None:
            SCHEME_B_MODEL_DIR.mkdir(parents=True, exist_ok=True)
            joblib.dump(self.models[-1], SCHEME_B_MODEL_FILE)
            joblib.dump(scaler, SCHEME_B_SCALER_FILE)
            with open(SCHEME_B_FEATURES_FILE, 'w', encoding='utf-8') as f:
                json.dump(feature_names, f, ensure_ascii=False)
            logger.info(f"已保存方案B模型: {SCHEME_B_MODEL_FILE}, scaler, {len(feature_names)} 个特征")

        # 6. 输出特征重要性
        top_feat = self.get_top_features(15)
        if not top_feat.empty:
            logger.info("\nTop 15 重要特征:")
            for i, (feat, imp) in enumerate(top_feat.items()):
                logger.info(f"  {i+1}. {feat}: {imp:.4f}")

        result['feature_importance'] = self.feature_importance
        return result


def run_ai_strategy(model_type: str = 'lightgbm', period: int = None,
                    start_date: str = None, end_date: str = None) -> dict:
    """方案B快捷入口"""
    strategy = AIStrategy(model_type=model_type)
    return strategy.run(period, start_date, end_date)


def get_live_signal(period: int = None, model_type: str = 'lightgbm',
                    start_date: str = None, end_date: str = None) -> int:
    """
    使用已保存的方案B模型对当前最新K线生成信号（供实盘/模拟盘调用）。
    需先运行过 run_ai_strategy 或 AIStrategy().run() 以生成 models/scheme_b_*.joblib 与 features.json。
    :param period: K线主周期，默认 15
    :param model_type: 未使用（模型已固定为磁盘上的 scheme_b_last_model.joblib）
    :param start_date, end_date: 可选，限制 load_kline 范围
    :return: 0=观望, 1=做多, 2=做空
    """
    if period is None:
        period = PRIMARY_PERIOD
    if not SCHEME_B_MODEL_FILE.exists() or not SCHEME_B_SCALER_FILE.exists() or not SCHEME_B_FEATURES_FILE.exists():
        logger.warning("方案B模型未找到，请先运行 AI 策略训练保存模型")
        return 0
    model = joblib.load(SCHEME_B_MODEL_FILE)
    scaler = joblib.load(SCHEME_B_SCALER_FILE)
    with open(SCHEME_B_FEATURES_FILE, 'r', encoding='utf-8') as f:
        feature_cols = json.load(f)
    X_last = get_latest_feature_row(period, feature_cols, start_date, end_date)
    if X_last.empty:
        return 0
    X_scaled = scaler.transform(X_last)
    X_scaled_df = pd.DataFrame(X_scaled, columns=feature_cols, index=X_last.index)
    proba = model.predict_proba(X_scaled_df)[0]  # (p0, p1, p2)
    pred = model.predict(X_scaled_df)[0]
    confidence_threshold = MC.get('confidence_threshold', 0.45)
    logger.info(f"方案B 预测概率: 观望={proba[0]:.2f} 做多={proba[1]:.2f} 做空={proba[2]:.2f} -> {int(pred)}")
    if proba.max() < confidence_threshold:
        logger.info(f"置信度 {proba.max():.2f} < {confidence_threshold}，视为观望")
        return 0
    return int(pred)