"""N-BEATS + DLinear 模型集成预测 - 交集策略"""

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

import torch
import pickle
import numpy as np
import pandas as pd
import pandas_ta as ta
from sklearn.preprocessing import StandardScaler
from sklearn.metrics import confusion_matrix
from data.fetcher import BinanceDataFetcher


# ==================== 预处理器定义 ====================

class ClassificationPreprocessor:
    """DLinear分类预处理器"""
    def __init__(self, lookback=168, forecast_horizon=1):
        self.lookback = lookback
        self.forecast_horizon = forecast_horizon
        self.scaler = StandardScaler()
        self.is_fitted = False

    def add_features(self, df):
        df = df.copy()
        df['rsi'] = ta.rsi(df['close'], length=14)
        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]
        return df

    def create_labels(self, df):
        df = df.copy()
        df['future_price'] = df['close'].shift(-self.forecast_horizon)
        df['future_return'] = (df['future_price'] - df['close']) / df['close']
        df['label'] = (df['future_return'] > 0).astype(int)
        return df[:-self.forecast_horizon]

    def preprocess(self, df, fit=True):
        df = self.add_features(df)
        df = self.create_labels(df)
        df = df.dropna().reset_index(drop=True)

        feature_cols = ['open', 'high', 'low', 'close', 'volume',
                       'rsi', 'macd', 'macd_hist', 'macd_signal']

        if fit:
            scaled_features = self.scaler.fit_transform(df[feature_cols])
            self.is_fitted = True
            self.feature_columns = feature_cols
        else:
            scaled_features = self.scaler.transform(df[feature_cols])

        X, y = [], []
        for i in range(len(scaled_features) - self.lookback):
            X.append(scaled_features[i:i + self.lookback])
            y.append(df['label'].iloc[i + self.lookback])

        return np.array(X), np.array(y), df


class NBeatsPreprocessor:
    """N-BEATS预处理器"""
    def __init__(self, lookback=168, forecast_horizon=1):
        self.lookback = lookback
        self.forecast_horizon = forecast_horizon
        self.price_scaler = StandardScaler()
        self.is_fitted = False

    def preprocess(self, df, fit=True):
        df = df.copy()
        df['future_price'] = df['close'].shift(-self.forecast_horizon)
        df = df[:-self.forecast_horizon].reset_index(drop=True)

        prices = df['close'].values.reshape(-1, 1)
        future_prices = df['future_price'].values.reshape(-1, 1)

        if fit:
            scaled_prices = self.price_scaler.fit_transform(prices)
            scaled_future = self.price_scaler.transform(future_prices)
            self.is_fitted = True
        else:
            scaled_prices = self.price_scaler.transform(prices)
            scaled_future = self.price_scaler.transform(future_prices)

        X, y = [], []
        for i in range(len(scaled_prices) - self.lookback):
            X.append(scaled_prices[i:i + self.lookback, 0])
            y.append(scaled_future[i + self.lookback, 0])

        return np.array(X), np.array(y), df

    def inverse_transform_price(self, scaled_price):
        return self.price_scaler.inverse_transform(scaled_price.reshape(-1, 1)).flatten()


def main():
    print("="*80)
    print("模型集成预测：N-BEATS + DLinear 交集策略")
    print("="*80)

    # ========== 步骤1: 加载两个模型 ==========
    print("\n[1/5] 加载两个模型...")

    device = torch.device('mps' if torch.backends.mps.is_available() else 'cpu')

    # 加载DLinear分类模型
    from models.dlinear.model import DLinear
    dlinear_checkpoint = torch.load(
        './models/checkpoints/dlinear_classification_2026/best_model.pt',
        map_location=device
    )
    dlinear_model = DLinear(dlinear_checkpoint['config'])
    dlinear_model.load_state_dict(dlinear_checkpoint['model_state_dict'])
    dlinear_model = dlinear_model.to(device)
    dlinear_model.eval()

    with open('./models/checkpoints/dlinear_classification_2026/preprocessor.pkl', 'rb') as f:
        dlinear_preprocessor = pickle.load(f)

    print(f"✓ DLinear分类模型已加载 (准确率: 53.10%)")

    # 加载N-BEATS回归模型
    from models.nbeats.model import NBeats
    nbeats_checkpoint = torch.load(
        './models/checkpoints/nbeats_btc_1h/best_model.pt',
        map_location=device
    )
    nbeats_model = NBeats(nbeats_checkpoint['config'])
    nbeats_model.load_state_dict(nbeats_checkpoint['model_state_dict'])
    nbeats_model = nbeats_model.to(device)
    nbeats_model.eval()

    with open('./models/checkpoints/nbeats_btc_1h/preprocessor.pkl', 'rb') as f:
        nbeats_preprocessor = pickle.load(f)

    print(f"✓ N-BEATS回归模型已加载 (准确率: 50.83%)")

    # ========== 步骤2: 获取2026年测试数据 ==========
    print("\n[2/5] 获取2026年测试数据...")
    fetcher = BinanceDataFetcher()

    try:
        df_2026 = fetcher.fetch_ohlcv('BTC', '1h', start_date='2026-01-01',
                                      end_date='2026-12-31', use_cache=False)
    except:
        df_2026 = fetcher.fetch_ohlcv('BTC', '1h', start_date='2025-12-01', use_cache=False)

    print(f"✓ 测试数据: {len(df_2026)} 行")
    print(f"✓ 时间: {df_2026['timestamp'].min()} 至 {df_2026['timestamp'].max()}")

    # ========== 步骤3: 预处理数据 ==========
    print("\n[3/5] 预处理数据...")

    # DLinear预处理
    X_dlinear, y_dlinear, df_dlinear = dlinear_preprocessor.preprocess(df_2026, fit=False)

    # N-BEATS预处理
    X_nbeats, y_nbeats, df_nbeats = nbeats_preprocessor.preprocess(df_2026, fit=False)

    # 确保两个模型的样本数量一致
    min_samples = min(len(X_dlinear), len(X_nbeats))
    X_dlinear = X_dlinear[:min_samples]
    y_dlinear = y_dlinear[:min_samples]
    X_nbeats = X_nbeats[:min_samples]
    y_nbeats = y_nbeats[:min_samples]

    print(f"✓ 对齐样本数: {min_samples}")

    # ========== 步骤4: 两个模型分别预测 ==========
    print("\n[4/5] 两个模型分别预测...")

    # DLinear预测
    X_dlinear_tensor = torch.FloatTensor(X_dlinear).to(device)
    with torch.no_grad():
        dlinear_logits = dlinear_model(X_dlinear_tensor).cpu().numpy()
        dlinear_probs = 1 / (1 + np.exp(-dlinear_logits))  # sigmoid
        dlinear_preds = (dlinear_probs > 0.5).astype(int).flatten()  # 1=涨, 0=跌

    print(f"✓ DLinear预测完成")
    print(f"  预测上涨: {(dlinear_preds==1).sum()} ({(dlinear_preds==1).sum()/len(dlinear_preds)*100:.1f}%)")
    print(f"  预测下跌: {(dlinear_preds==0).sum()} ({(dlinear_preds==0).sum()/len(dlinear_preds)*100:.1f}%)")

    # N-BEATS预测
    X_nbeats_tensor = torch.FloatTensor(X_nbeats).to(device)
    with torch.no_grad():
        nbeats_pred_scaled = nbeats_model(X_nbeats_tensor).cpu().numpy()

    # 还原真实价格
    nbeats_pred_prices = nbeats_preprocessor.inverse_transform_price(nbeats_pred_scaled)

    # 获取当前价格（用于判断方向）
    current_prices = df_nbeats['close'].iloc[168:168+len(nbeats_pred_prices)].values
    nbeats_preds = (nbeats_pred_prices > current_prices).astype(int)  # 1=涨, 0=跌

    print(f"✓ N-BEATS预测完成")
    print(f"  预测上涨: {(nbeats_preds==1).sum()} ({(nbeats_preds==1).sum()/len(nbeats_preds)*100:.1f}%)")
    print(f"  预测下跌: {(nbeats_preds==0).sum()} ({(nbeats_preds==0).sum()/len(nbeats_preds)*100:.1f}%)")

    # ========== 步骤5: 交集策略 ==========
    print("\n[5/5] 应用交集策略...")

    # 真实标签
    true_labels = y_dlinear

    # 策略1: 两个模型都预测上涨 → 上涨
    # 策略2: 两个模型都预测下跌 → 下跌
    # 策略3: 两个模型预测不一致 → 不做决策（弃权）

    both_agree_up = (dlinear_preds == 1) & (nbeats_preds == 1)
    both_agree_down = (dlinear_preds == 0) & (nbeats_preds == 0)
    disagree = (dlinear_preds != nbeats_preds)

    n_agree_up = both_agree_up.sum()
    n_agree_down = both_agree_down.sum()
    n_disagree = disagree.sum()

    print(f"\n📊 预测一致性分析:")
    print(f"  双方都预测上涨: {n_agree_up} ({n_agree_up/len(dlinear_preds)*100:.1f}%)")
    print(f"  双方都预测下跌: {n_agree_down} ({n_agree_down/len(dlinear_preds)*100:.1f}%)")
    print(f"  双方预测不一致: {n_disagree} ({n_disagree/len(dlinear_preds)*100:.1f}%)")

    # 只在两个模型一致时做决策
    agree_mask = both_agree_up | both_agree_down
    ensemble_preds = np.full(len(dlinear_preds), -1)  # -1表示弃权
    ensemble_preds[both_agree_up] = 1
    ensemble_preds[both_agree_down] = 0

    # 有效预测（非弃权）
    valid_mask = ensemble_preds != -1
    valid_preds = ensemble_preds[valid_mask]
    valid_labels = true_labels[valid_mask]

    print(f"\n📊 集成策略结果:")
    print(f"  有效预测数: {valid_mask.sum()} ({valid_mask.sum()/len(dlinear_preds)*100:.1f}%)")
    print(f"  弃权数: {n_disagree} ({n_disagree/len(dlinear_preds)*100:.1f}%)")

    if len(valid_preds) > 0:
        ensemble_accuracy = (valid_preds == valid_labels).mean()
        print(f"\n✨ 交集策略准确率: {ensemble_accuracy*100:.2f}%")

        # 混淆矩阵
        cm = confusion_matrix(valid_labels, valid_preds)
        tn, fp, fn, tp = cm.ravel()

        print(f"\n详细分析（仅有效预测）:")
        print(f"  ✓ 预测涨且真涨: {tp} ({tp/len(valid_preds)*100:.1f}%)")
        print(f"  ✓ 预测跌且真跌: {tn} ({tn/len(valid_preds)*100:.1f}%)")
        print(f"  ✗ 预测涨但真跌: {fp} ({fp/len(valid_preds)*100:.1f}%)")
        print(f"  ✗ 预测跌但真涨: {fn} ({fn/len(valid_preds)*100:.1f}%)")

    # ========== 对比分析 ==========
    print("\n" + "="*80)
    print("📊 性能对比")
    print("="*80)

    dlinear_acc = (dlinear_preds == true_labels).mean()
    nbeats_acc = (nbeats_preds == true_labels).mean()

    print(f"\n单模型性能:")
    print(f"  DLinear分类: {dlinear_acc*100:.2f}%")
    print(f"  N-BEATS回归: {nbeats_acc*100:.2f}%")

    if len(valid_preds) > 0:
        print(f"\n集成策略性能:")
        print(f"  交集准确率: {ensemble_accuracy*100:.2f}%")
        print(f"  覆盖率: {valid_mask.sum()/len(dlinear_preds)*100:.1f}%")
        print(f"  弃权率: {n_disagree/len(dlinear_preds)*100:.1f}%")

        # 性能评估
        print("\n" + "="*80)
        print("🎯 策略评估")
        print("="*80)

        if ensemble_accuracy > max(dlinear_acc, nbeats_acc):
            improvement = ensemble_accuracy - max(dlinear_acc, nbeats_acc)
            print(f"\n  ✅ 交集策略优于单模型!")
            print(f"     准确率提升: {improvement*100:.2f}个百分点")
            print(f"     代价: 弃权{n_disagree}次预测 ({n_disagree/len(dlinear_preds)*100:.1f}%)")
        else:
            print(f"\n  ⚠️  交集策略未能提升准确率")
            print(f"     可能原因: 两个模型一致时也可能都错")

        # 计算如果弃权时不交易的盈利
        print(f"\n💰 实际收益分析:")
        print(f"  如果只在两个模型都同意时交易:")
        print(f"    - 交易次数: {valid_mask.sum()} ({valid_mask.sum()/len(dlinear_preds)*100:.1f}%)")
        print(f"    - 准确率: {ensemble_accuracy*100:.2f}%")
        print(f"    - 避免交易: {n_disagree}次 (规避不确定性)")

    print("\n" + "="*80)
    print("✓ 集成预测分析完成!")
    print("="*80)


if __name__ == "__main__":
    main()