"""四模型集成每日准确率分析 - 多时间段测试（2019年1月、2020年1月、2026年1月）"""

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 predict_period(period_name, start_date, end_date, dlinear_model, nbeats_model, ttm_model, timesfm_model,
                   dlinear_preprocessor, nbeats_preprocessor, device):
    """预测单个时间段的数据"""

    print(f"\n{'='*80}")
    print(f"📅 {period_name} ({start_date} 至 {end_date})")
    print(f"{'='*80}")

    # 获取数据
    fetcher = BinanceDataFetcher()
    df = fetcher.fetch_ohlcv('BTC', '1h', start_date=start_date, end_date=end_date, use_cache=False)

    print(f"✓ 数据: {len(df)} 行")

    # 预处理
    X_dlinear, y_dlinear, df_dlinear = dlinear_preprocessor.preprocess(df, fit=False)
    X_nbeats, y_nbeats, df_nbeats = nbeats_preprocessor.preprocess(df, 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]

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

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

    # 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['close'].values
    all_open_prices = df['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)

    # TimesFM预测
    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)

    # 集成策略
    true_labels = y_dlinear

    # 双模型交集
    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

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

    # 四模型投票
    vote_sum_4 = dlinear_preds + nbeats_preds + ttm_preds + timesfm_preds
    four_model_preds = np.where(vote_sum_4 >= 3, 1, 0)

    # 构建每日结果DataFrame
    timestamps = df['timestamp'].iloc[168:168+min_samples].values
    timestamps_utc = pd.to_datetime(timestamps)
    timestamps_et = timestamps_utc.tz_localize('UTC').tz_convert('America/New_York')
    dates = timestamps_et.date

    results_df = pd.DataFrame({
        'date': dates,
        'timestamp_et': timestamps_et,
        'true_label': true_labels,
        'dlinear_pred': dlinear_preds,
        'nbeats_pred': nbeats_preds,
        'ttm_pred': ttm_preds,
        'timesfm_pred': timesfm_preds,
        'two_model_pred': two_model_preds,
        'three_model_pred': three_model_preds,
        'four_model_pred': four_model_preds,
    })

    # 按天统计
    daily_stats = results_df.groupby('date').agg({
        'true_label': 'count',
        'dlinear_pred': lambda x: (x == results_df.loc[x.index, 'true_label']).sum(),
        'nbeats_pred': lambda x: (x == results_df.loc[x.index, 'true_label']).sum(),
        'ttm_pred': lambda x: (x == results_df.loc[x.index, 'true_label']).sum(),
        'timesfm_pred': lambda x: (x == results_df.loc[x.index, 'true_label']).sum(),
        'two_model_pred': lambda x: ((x == results_df.loc[x.index, 'true_label']) & (x != -1)).sum(),
        'three_model_pred': lambda x: (x == results_df.loc[x.index, 'true_label']).sum(),
        'four_model_pred': lambda x: (x == results_df.loc[x.index, 'true_label']).sum(),
    })

    daily_stats.columns = ['total', 'dlinear_correct', 'nbeats_correct', 'ttm_correct',
                            'timesfm_correct', 'two_model_correct', 'three_model_correct', 'four_model_correct']

    # 计算每天的有效预测数（双模型）
    two_model_valid = results_df.groupby('date')['two_model_pred'].apply(lambda x: (x != -1).sum())
    daily_stats['two_model_valid'] = two_model_valid

    # 计算准确率
    daily_stats['dlinear_acc'] = daily_stats['dlinear_correct'] / daily_stats['total'] * 100
    daily_stats['nbeats_acc'] = daily_stats['nbeats_correct'] / daily_stats['total'] * 100
    daily_stats['ttm_acc'] = daily_stats['ttm_correct'] / daily_stats['total'] * 100
    daily_stats['timesfm_acc'] = daily_stats['timesfm_correct'] / daily_stats['total'] * 100
    daily_stats['two_model_acc'] = daily_stats['two_model_correct'] / daily_stats['two_model_valid'] * 100
    daily_stats['three_model_acc'] = daily_stats['three_model_correct'] / daily_stats['total'] * 100
    daily_stats['four_model_acc'] = daily_stats['four_model_correct'] / daily_stats['total'] * 100

    daily_stats['two_model_acc'] = daily_stats['two_model_acc'].fillna(0)

    return daily_stats, results_df


def main():
    print("="*80)
    print("四模型集成每日准确率分析 - 多时间段测试")
    print("测试时段: 2019年1月、2020年1月、2026年1月")
    print("="*80)

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

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

    from models.dlinear.model import DLinear
    dlinear_checkpoint = torch.load(
        './models/checkpoints/dlinear_classification_2026/best_model.pt',
        map_location=device, weights_only=False
    )
    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)

    from models.nbeats.model import NBeats
    nbeats_checkpoint = torch.load(
        './models/checkpoints/nbeats_btc_1h/best_model.pt',
        map_location=device, weights_only=False
    )
    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)

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

    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("✓ 四个模型已加载 (DLinear, N-BEATS, TTM-R2, TimesFM)")

    # ========== 多时间段预测 ==========
    print("\n[2/2] 多时间段预测...")

    periods = [
        ("2019年1月", "2019-01-01", "2019-01-31"),
        ("2020年1月", "2020-01-01", "2020-01-31"),
        ("2026年1月", "2026-01-01", "2026-01-31"),
    ]

    all_period_stats = {}

    for period_name, start_date, end_date in periods:
        daily_stats, results_df = predict_period(
            period_name, start_date, end_date,
            dlinear_model, nbeats_model, ttm_model, timesfm_model,
            dlinear_preprocessor, nbeats_preprocessor, device
        )
        all_period_stats[period_name] = daily_stats

        # 显示每日准确率
        print(f"\n每日准确率明细:")
        print(f"{'日期':<12} {'样本':>4} {'DLinear':>8} {'N-BEATS':>8} {'TTM-R2':>8} {'TimesFM':>8} {'双模型':>8} {'三模型':>8} {'四模型':>8}")
        print("-" * 95)

        for date, row in daily_stats.iterrows():
            print(f"{date} {row['total']:>4.0f} "
                  f"{row['dlinear_acc']:>7.1f}% "
                  f"{row['nbeats_acc']:>7.1f}% "
                  f"{row['ttm_acc']:>7.1f}% "
                  f"{row['timesfm_acc']:>7.1f}% "
                  f"{row['two_model_acc']:>7.1f}% "
                  f"{row['three_model_acc']:>7.1f}% "
                  f"{row['four_model_acc']:>7.1f}%")

        # 统计总体
        print(f"\n{period_name} 总体统计:")
        print(f"  DLinear:  {daily_stats['dlinear_correct'].sum() / daily_stats['total'].sum() * 100:.2f}%")
        print(f"  N-BEATS:  {daily_stats['nbeats_correct'].sum() / daily_stats['total'].sum() * 100:.2f}%")
        print(f"  TTM-R2:   {daily_stats['ttm_correct'].sum() / daily_stats['total'].sum() * 100:.2f}%")
        print(f"  TimesFM:  {daily_stats['timesfm_correct'].sum() / daily_stats['total'].sum() * 100:.2f}%")
        print(f"  双模型:   {daily_stats['two_model_correct'].sum() / daily_stats['two_model_valid'].sum() * 100:.2f}% (覆盖率: {daily_stats['two_model_valid'].sum() / daily_stats['total'].sum() * 100:.1f}%)")
        print(f"  三模型:   {daily_stats['three_model_correct'].sum() / daily_stats['total'].sum() * 100:.2f}%")
        print(f"  四模型:   {daily_stats['four_model_correct'].sum() / daily_stats['total'].sum() * 100:.2f}%")

    # ========== 跨时期对比 ==========
    print("\n" + "="*80)
    print("📊 跨时期对比总结")
    print("="*80)

    print(f"\n{'时期':<12} {'DLinear':>9} {'N-BEATS':>9} {'TTM-R2':>9} {'TimesFM':>9} {'双模型':>9} {'三模型':>9} {'四模型':>9}")
    print("-" * 95)

    for period_name in ["2019年1月", "2020年1月", "2026年1月"]:
        stats = all_period_stats[period_name]
        print(f"{period_name:<12} "
              f"{stats['dlinear_correct'].sum() / stats['total'].sum() * 100:>8.2f}% "
              f"{stats['nbeats_correct'].sum() / stats['total'].sum() * 100:>8.2f}% "
              f"{stats['ttm_correct'].sum() / stats['total'].sum() * 100:>8.2f}% "
              f"{stats['timesfm_correct'].sum() / stats['total'].sum() * 100:>8.2f}% "
              f"{stats['two_model_correct'].sum() / stats['two_model_valid'].sum() * 100:>8.2f}% "
              f"{stats['three_model_correct'].sum() / stats['total'].sum() * 100:>8.2f}% "
              f"{stats['four_model_correct'].sum() / stats['total'].sum() * 100:>8.2f}%")

    print("\n" + "="*80)
    print("✓ 多时间段每日准确率分析完成!")
    print("="*80)


if __name__ == "__main__":
    main()