"""五模型集成预测：DLinear + N-BEATS + TTM-R2 + TimesFM 2.5 + CatBoost"""

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 catboost import CatBoostClassifier
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 add_catboost_features(df):
    """为CatBoost添加丰富的技术指标特征"""
    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
    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
    df['obv'] = ta.obv(df['close'], df['volume'])

    # Stochastic
    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
    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 main():
    print("="*80)
    print("五模型集成预测：DLinear + N-BEATS + TTM-R2 + TimesFM + CatBoost")
    print("="*80)

    # ========== 步骤1: 加载PyTorch模型 ==========
    print("\n[1/8] 加载DLinear和N-BEATS模型...")

    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已加载")

    # 加载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已加载")

    # ========== 步骤2: 加载TTM-R2 ==========
    print("\n[2/8] 加载TTM-R2零样本模型...")

    from tsfm_public.models.tinytimemixer import TinyTimeMixerForPrediction

    ttm_model = TinyTimeMixerForPrediction.from_pretrained(
        "ibm-granite/granite-timeseries-ttm-r2",
        trust_remote_code=True
    )
    ttm_model = ttm_model.to(device)
    ttm_model.eval()

    print(f"✓ TTM-R2已加载")

    # ========== 步骤3: 加载TimesFM 2.5 ==========
    print("\n[3/8] 加载TimesFM 2.5模型...")

    import timesfm

    timesfm_model = timesfm.TimesFM_2p5_200M_torch.from_pretrained(
        "google/timesfm-2.5-200m-pytorch",
        torch_compile=False
    )

    timesfm_model.compile(
        timesfm.ForecastConfig(
            max_context=512,
            max_horizon=1,
            normalize_inputs=True,
            use_continuous_quantile_head=True,
            force_flip_invariance=True,
            infer_is_positive=True,
            fix_quantile_crossing=True,
        )
    )

    print(f"✓ TimesFM 2.5已加载")

    # ========== 步骤4: 加载CatBoost ==========
    print("\n[4/8] 加载CatBoost模型...")

    catboost_model = CatBoostClassifier()
    catboost_model.load_model('./models/checkpoints/catboost_classification_2026/catboost_model.cbm')

    with open('./models/checkpoints/catboost_classification_2026/metadata.pkl', 'rb') as f:
        catboost_metadata = pickle.load(f)
        catboost_feature_cols = catboost_metadata['feature_cols']

    print(f"✓ CatBoost已加载 (特征数: {len(catboost_feature_cols)})")

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

    fetcher = BinanceDataFetcher()
    df_2026 = fetcher.fetch_ohlcv('BTC', '1h', start_date='2026-01-01',
                                   end_date='2026-01-31', use_cache=False)

    print(f"✓ 数据获取完成: {len(df_2026)} 行")

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

    X_dlinear, y_dlinear, df_dlinear = dlinear_preprocessor.preprocess(df_2026, fit=False)
    X_nbeats, y_nbeats, df_nbeats = nbeats_preprocessor.preprocess(df_2026, fit=False)

    # CatBoost特征工程
    df_catboost = add_catboost_features(df_2026)
    df_catboost = df_catboost.dropna().reset_index(drop=True)
    X_catboost = df_catboost[catboost_feature_cols]

    # 确保样本数量一致
    min_samples = min(len(X_dlinear), len(X_nbeats), len(X_catboost))
    X_dlinear = X_dlinear[:min_samples]
    y_dlinear = y_dlinear[:min_samples]
    X_nbeats = X_nbeats[:min_samples]
    X_catboost = X_catboost.iloc[:min_samples]

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

    # ========== 步骤7: 五个模型分别预测 ==========
    print("\n[7/8] 五个模型分别预测...")

    # 7.1 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))
        dlinear_preds = (dlinear_probs > 0.5).astype(int).flatten()

    print(f"✓ DLinear: 涨{(dlinear_preds==1).sum()} ({(dlinear_preds==1).sum()/len(dlinear_preds)*100:.1f}%)")

    # 7.2 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_open_prices = df_nbeats['open'].iloc[168:168+len(nbeats_pred_prices)].values
    nbeats_preds = (nbeats_pred_prices > current_open_prices).astype(int)

    print(f"✓ N-BEATS: 涨{(nbeats_preds==1).sum()} ({(nbeats_preds==1).sum()/len(nbeats_preds)*100:.1f}%)")

    # 7.3 TTM-R2预测
    context_length = 512
    ttm_preds_list = []
    all_prices = df_2026['close'].values
    all_open_prices = df_2026['open'].values

    for i in range(min_samples):
        start_idx = i
        end_idx = start_idx + context_length

        if end_idx > len(all_prices):
            context_data = all_prices[-context_length:].reshape(1, -1, 1)
            current_open_price = all_open_prices[min(168 + i, len(all_open_prices) - 1)]
        else:
            context_data = all_prices[start_idx:end_idx].reshape(1, -1, 1)
            current_open_price = all_open_prices[end_idx] if end_idx < len(all_open_prices) else all_open_prices[-1]

        context_tensor = torch.FloatTensor(context_data).to(device)

        with torch.no_grad():
            output = ttm_model(context_tensor)
            pred_price = output.prediction_outputs.cpu().numpy()[0, 0, 0]

        ttm_preds_list.append(1 if pred_price > current_open_price else 0)

    ttm_preds = np.array(ttm_preds_list)

    print(f"✓ TTM-R2: 涨{(ttm_preds==1).sum()} ({(ttm_preds==1).sum()/len(ttm_preds)*100:.1f}%)")

    # 7.4 TimesFM 2.5预测
    timesfm_preds_list = []

    for i in range(min_samples):
        start_idx = i
        end_idx = start_idx + 512

        if end_idx > len(all_prices):
            context_prices = all_prices[-512:]
            current_open_price = all_open_prices[min(168 + i, len(all_open_prices) - 1)]
        else:
            context_prices = all_prices[start_idx:end_idx]
            current_open_price = all_open_prices[end_idx] if end_idx < len(all_open_prices) else all_open_prices[-1]

        try:
            point_forecast, _ = timesfm_model.forecast(
                horizon=1,
                inputs=[context_prices],
            )
            pred_price = point_forecast[0, 0]
            timesfm_preds_list.append(1 if pred_price > current_open_price else 0)
        except:
            timesfm_preds_list.append(np.random.randint(0, 2))

    timesfm_preds = np.array(timesfm_preds_list)

    print(f"✓ TimesFM: 涨{(timesfm_preds==1).sum()} ({(timesfm_preds==1).sum()/len(timesfm_preds)*100:.1f}%)")

    # 7.5 CatBoost预测
    catboost_preds = catboost_model.predict(X_catboost).astype(int)

    print(f"✓ CatBoost: 涨{(catboost_preds==1).sum()} ({(catboost_preds==1).sum()/len(catboost_preds)*100:.1f}%)")

    # ========== 步骤8: 集成策略评估 ==========
    print("\n[8/8] 集成策略评估...")

    true_labels = y_dlinear

    # 策略1: 双模型交集（DLinear + N-BEATS）
    two_model_agree_up = (dlinear_preds == 1) & (nbeats_preds == 1)
    two_model_agree_down = (dlinear_preds == 0) & (nbeats_preds == 0)
    two_model_preds = np.full(len(dlinear_preds), -1)
    two_model_preds[two_model_agree_up] = 1
    two_model_preds[two_model_agree_down] = 0

    two_model_valid = two_model_preds != -1
    two_model_correct = (two_model_preds == true_labels) & two_model_valid
    two_model_acc = two_model_correct.sum() / two_model_valid.sum() * 100 if two_model_valid.sum() > 0 else 0
    two_model_coverage = two_model_valid.sum() / len(true_labels) * 100

    # 策略2: 三模型投票（DLinear + N-BEATS + TTM-R2）
    vote_sum_3 = dlinear_preds + nbeats_preds + ttm_preds
    three_model_preds = np.where(vote_sum_3 >= 2, 1, 0)
    three_model_acc = (three_model_preds == true_labels).mean() * 100

    # 策略3: 四模型投票
    vote_sum_4 = dlinear_preds + nbeats_preds + ttm_preds + timesfm_preds
    four_model_preds = np.where(vote_sum_4 >= 3, 1, 0)
    four_model_acc = (four_model_preds == true_labels).mean() * 100

    # 策略4: 五模型强投票（≥3票）
    vote_sum_5 = dlinear_preds + nbeats_preds + ttm_preds + timesfm_preds + catboost_preds
    five_model_preds_strong = np.where(vote_sum_5 >= 3, 1, 0)
    five_model_acc_strong = (five_model_preds_strong == true_labels).mean() * 100

    # 策略5: 五模型弱投票（≥2票）
    five_model_preds_weak = np.where(vote_sum_5 >= 2, 1, 0)
    five_model_acc_weak = (five_model_preds_weak == true_labels).mean() * 100

    # ========== 结果展示 ==========
    print("\n" + "="*80)
    print("📊 五模型集成策略对比结果")
    print("="*80)

    print(f"\n单模型性能：")
    print(f"  DLinear:     {(dlinear_preds == true_labels).mean() * 100:.2f}%")
    print(f"  CatBoost:    {(catboost_preds == true_labels).mean() * 100:.2f}%")
    print(f"  TimesFM 2.5: {(timesfm_preds == true_labels).mean() * 100:.2f}%")
    print(f"  TTM-R2:      {(ttm_preds == true_labels).mean() * 100:.2f}%")
    print(f"  N-BEATS:     {(nbeats_preds == true_labels).mean() * 100:.2f}%")

    print(f"\n集成策略性能：")
    print(f"  双模型交集 (DL+NB):            {two_model_acc:.2f}% (覆盖率: {two_model_coverage:.2f}%)")
    print(f"  三模型投票 (DL+NB+TTM, ≥2票):  {three_model_acc:.2f}% (覆盖率: 100%)")
    print(f"  四模型投票 (前4, ≥3票):        {four_model_acc:.2f}% (覆盖率: 100%)")
    print(f"  五模型强投票 (所有, ≥3票):     {five_model_acc_strong:.2f}% (覆盖率: 100%)")
    print(f"  五模型弱投票 (所有, ≥2票):     {five_model_acc_weak:.2f}% (覆盖率: 100%)")

    # 投票分布分析
    print(f"\n五模型投票分布：")
    for votes in range(6):
        count = (vote_sum_5 == votes).sum()
        pct = count / len(vote_sum_5) * 100
        bars = "█" * int(pct / 2)
        print(f"  {votes}票同意: {count:3d}次 ({pct:5.1f}%) {bars}")

    # 模型一致性分析
    print(f"\n模型一致性分析：")
    all_agree_up = (dlinear_preds == 1) & (nbeats_preds == 1) & (ttm_preds == 1) & (timesfm_preds == 1) & (catboost_preds == 1)
    all_agree_down = (dlinear_preds == 0) & (nbeats_preds == 0) & (ttm_preds == 0) & (timesfm_preds == 0) & (catboost_preds == 0)
    all_agree = all_agree_up | all_agree_down

    print(f"  五模型完全一致: {all_agree.sum()}次 ({all_agree.sum()/len(all_agree)*100:.1f}%)")
    if all_agree.sum() > 0:
        all_agree_acc = ((five_model_preds_strong[all_agree] == true_labels[all_agree]).mean() * 100)
        print(f"    一致时准确率: {all_agree_acc:.2f}%")

    print("\n" + "="*80)
    print("✓ 五模型集成预测完成!")
    print("="*80)


if __name__ == "__main__":
    main()