"""分析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 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("DLinear输出概率值分布分析 - 以0.1阶梯统计")
    print("="*80)

    # ========== 加载模型 ==========
    print("\n[1/4] 加载模型...")

    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)

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

    # 加载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"✓ 三个模型已加载")

    # ========== 获取数据 ==========
    print("\n[2/4] 获取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("\n[3/4] 预处理数据...")

    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)

    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]

    # ========== 预测并获取DLinear概率值 ==========
    print("\n[4/4] 执行预测并分析DLinear输出概率...")

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

    # 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)

    # 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)

    # 双模型交集
    two_model_agree_up = (dlinear_preds == 1) & (nbeats_preds == 1)
    two_model_agree_down = (dlinear_preds == 0) & (nbeats_preds == 0)
    two_model_valid = two_model_agree_up | two_model_agree_down

    # 三模型多数投票
    vote_sum = dlinear_preds + nbeats_preds + ttm_preds
    three_model_preds = np.where(vote_sum >= 2, 1, 0)

    # 真实标签
    true_labels = y_dlinear

    print(f"✓ 预测完成，开始分析DLinear概率分布...")

    # ========== 按0.1阶梯分析 ==========
    print("\n" + "="*80)
    print("📊 DLinear输出概率分布分析（0.1阶梯）")
    print("="*80)

    # 定义概率区间
    prob_bins = np.arange(0, 1.1, 0.1)
    prob_labels = [f"{i/10:.1f}-{(i+1)/10:.1f}" for i in range(10)]

    # 为每个样本分配到概率区间
    prob_bin_indices = np.digitize(dlinear_probs, prob_bins) - 1
    prob_bin_indices = np.clip(prob_bin_indices, 0, 9)  # 确保在0-9范围内

    # 初始化统计
    stats = []

    for i, bin_label in enumerate(prob_labels):
        mask = prob_bin_indices == i
        count = mask.sum()

        if count == 0:
            continue

        # 该区间的真实标签
        bin_true = true_labels[mask]

        # 该区间DLinear的预测
        bin_dlinear_pred = dlinear_preds[mask]

        # 该区间在双模型交集中的表现
        bin_two_model_valid = two_model_valid[mask]
        bin_two_model_pred = np.full(count, -1)
        bin_two_model_pred[two_model_agree_up[mask]] = 1
        bin_two_model_pred[two_model_agree_down[mask]] = 0

        # 该区间在三模型投票中的表现
        bin_three_model_pred = three_model_preds[mask]

        # 计算统计
        dlinear_acc = (bin_dlinear_pred == bin_true).mean() * 100 if count > 0 else 0

        two_model_valid_count = bin_two_model_valid.sum()
        two_model_acc = ((bin_two_model_pred == bin_true) & (bin_two_model_pred != -1)).sum() / two_model_valid_count * 100 if two_model_valid_count > 0 else 0
        two_model_coverage = two_model_valid_count / count * 100

        three_model_acc = (bin_three_model_pred == bin_true).mean() * 100 if count > 0 else 0

        # 平均概率值
        avg_prob = dlinear_probs[mask].mean()

        # 预测倾向
        pred_up = bin_dlinear_pred.mean() * 100

        stats.append({
            'bin': bin_label,
            'count': count,
            'avg_prob': avg_prob,
            'pred_up_pct': pred_up,
            'dlinear_acc': dlinear_acc,
            'two_model_acc': two_model_acc,
            'two_model_cov': two_model_coverage,
            'three_model_acc': three_model_acc
        })

    # 转换为DataFrame
    stats_df = pd.DataFrame(stats)

    # 显示统计表
    print("\n概率区间分析：")
    print(f"{'区间':<10} {'样本数':>6} {'平均概率':>8} {'预测↑%':>7} "
          f"{'DL准确率':>9} {'双模准确率':>10} {'双模覆盖':>8} {'三模准确率':>10}")
    print("-" * 95)

    for _, row in stats_df.iterrows():
        # 根据准确率添加标记
        dl_mark = "🔥" if row['dlinear_acc'] >= 70 else ("✅" if row['dlinear_acc'] >= 60 else ("⚠️" if row['dlinear_acc'] < 45 else ""))
        two_mark = "🔥" if row['two_model_acc'] >= 70 else ("✅" if row['two_model_acc'] >= 60 else ("⚠️" if row['two_model_acc'] < 45 else ""))
        three_mark = "🔥" if row['three_model_acc'] >= 70 else ("✅" if row['three_model_acc'] >= 60 else ("⚠️" if row['three_model_acc'] < 45 else ""))

        print(f"{row['bin']:<10} {row['count']:>6} {row['avg_prob']:>8.3f} {row['pred_up_pct']:>6.1f}% "
              f"{row['dlinear_acc']:>8.1f}{dl_mark} "
              f"{row['two_model_acc']:>9.1f}{two_mark} "
              f"{row['two_model_cov']:>7.1f}% "
              f"{row['three_model_acc']:>9.1f}{three_mark}")

    # ========== 关键发现 ==========
    print("\n" + "="*80)
    print("🔍 关键发现")
    print("="*80)

    # 极端概率区间（高置信度）
    high_conf_up = stats_df[stats_df['avg_prob'] >= 0.65]
    high_conf_down = stats_df[stats_df['avg_prob'] <= 0.35]

    print(f"\n📈 高置信度看涨区间 (概率≥0.65):")
    if len(high_conf_up) > 0:
        total_samples = high_conf_up['count'].sum()
        avg_dl_acc = (high_conf_up['dlinear_acc'] * high_conf_up['count']).sum() / total_samples
        avg_two_acc = (high_conf_up['two_model_acc'] * high_conf_up['count']).sum() / total_samples
        avg_three_acc = (high_conf_up['three_model_acc'] * high_conf_up['count']).sum() / total_samples

        print(f"  样本数: {total_samples} ({total_samples/len(dlinear_probs)*100:.1f}%)")
        print(f"  DLinear平均准确率: {avg_dl_acc:.2f}%")
        print(f"  双模型平均准确率: {avg_two_acc:.2f}%")
        print(f"  三模型平均准确率: {avg_three_acc:.2f}%")
    else:
        print("  无数据")

    print(f"\n📉 高置信度看跌区间 (概率≤0.35):")
    if len(high_conf_down) > 0:
        total_samples = high_conf_down['count'].sum()
        avg_dl_acc = (high_conf_down['dlinear_acc'] * high_conf_down['count']).sum() / total_samples
        avg_two_acc = (high_conf_down['two_model_acc'] * high_conf_down['count']).sum() / total_samples
        avg_three_acc = (high_conf_down['three_model_acc'] * high_conf_down['count']).sum() / total_samples

        print(f"  样本数: {total_samples} ({total_samples/len(dlinear_probs)*100:.1f}%)")
        print(f"  DLinear平均准确率: {avg_dl_acc:.2f}%")
        print(f"  双模型平均准确率: {avg_two_acc:.2f}%")
        print(f"  三模型平均准确率: {avg_three_acc:.2f}%")
    else:
        print("  无数据")

    # 中等置信度区间
    mid_conf = stats_df[(stats_df['avg_prob'] > 0.45) & (stats_df['avg_prob'] < 0.55)]

    print(f"\n⚖️  中等置信度区间 (0.45-0.55):")
    if len(mid_conf) > 0:
        total_samples = mid_conf['count'].sum()
        avg_dl_acc = (mid_conf['dlinear_acc'] * mid_conf['count']).sum() / total_samples
        avg_two_acc = (mid_conf['two_model_acc'] * mid_conf['count']).sum() / total_samples if total_samples > 0 else 0
        avg_three_acc = (mid_conf['three_model_acc'] * mid_conf['count']).sum() / total_samples

        print(f"  样本数: {total_samples} ({total_samples/len(dlinear_probs)*100:.1f}%)")
        print(f"  DLinear平均准确率: {avg_dl_acc:.2f}%")
        print(f"  双模型平均准确率: {avg_two_acc:.2f}%")
        print(f"  三模型平均准确率: {avg_three_acc:.2f}%")
        print(f"  ⚠️  建议: 低置信度区间，考虑弃权或减仓")
    else:
        print("  无数据")

    # 找出表现最好的概率区间
    best_bin = stats_df.loc[stats_df['dlinear_acc'].idxmax()]
    worst_bin = stats_df.loc[stats_df['dlinear_acc'].idxmin()]

    print(f"\n🏆 DLinear表现最佳区间:")
    print(f"  区间: {best_bin['bin']}")
    print(f"  准确率: {best_bin['dlinear_acc']:.1f}%")
    print(f"  样本数: {best_bin['count']}")

    print(f"\n💀 DLinear表现最差区间:")
    print(f"  区间: {worst_bin['bin']}")
    print(f"  准确率: {worst_bin['dlinear_acc']:.1f}%")
    print(f"  样本数: {worst_bin['count']}")

    # ========== 概率分布可视化 ==========
    print("\n" + "="*80)
    print("📊 DLinear概率分布直方图")
    print("="*80)

    print("\n样本数分布:")
    max_count = stats_df['count'].max()
    for _, row in stats_df.iterrows():
        bar_length = int(row['count'] / max_count * 40)
        bar = "█" * bar_length
        print(f"{row['bin']:<10} {bar} {row['count']}")

    print("\n" + "="*80)
    print("✓ DLinear概率分布分析完成!")
    print("="*80)


if __name__ == "__main__":
    main()