"""2026年2月BTC数据 - DLinear, N-BEATS, TTM-R2, DeepSeek-R1(高并发版) 准确率测试

输入数据：预测K线的前一根（即用截止到t的数据预测t+1）
价格对比：预测结果与预测K线(t+1)的收盘价做对比
时间：币安数据UTC，输出转为北京时间(UTC+8)，按北京时间分天
输出：每小时详细数据、每天准确率、总准确率、四模型集成、美股开盘标记
"""

import sys
import os
import json
import re
import time
import asyncio
import aiohttp
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

import torch
import numpy as np
import pandas as pd
import pickle
import pandas_ta as ta
from datetime import timedelta
from sklearn.preprocessing import StandardScaler
from data.fetcher import BinanceDataFetcher
from tsfm_public.models.tinytimemixer import TinyTimeMixerForPrediction
from utils.network_proxy import apply_global_proxy_from_env

apply_global_proxy_from_env()

# ==================== 时区与美股判断 ====================

def utc_to_beijing(ts):
    return ts + timedelta(hours=8)

def is_us_stock_open_by_utc(utc_ts):
    et_ts = utc_ts - timedelta(hours=5)
    et_date = pd.Timestamp(et_ts.date())
    us_holidays_feb_2026 = {pd.Timestamp('2026-02-16')}
    if et_date.weekday() >= 5 or et_date in us_holidays_feb_2026:
        return False
    return True

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

class ClassificationPreprocessor:
    def __init__(self, lookback=168, forecast_horizon=1):
        self.lookback = lookback
        self.forecast_horizon = forecast_horizon
        self.scaler = StandardScaler()
        self.is_fitted = False

class NBeatsPreprocessor:
    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()
        prices = df['close'].values.reshape(-1, 1)
        if fit:
            scaled_prices = self.price_scaler.fit_transform(prices)
            self.is_fitted = True
        else:
            scaled_prices = self.price_scaler.transform(prices)
        X, valid_indices = [], []
        for i in range(len(scaled_prices) - self.lookback):
            X.append(scaled_prices[i:i + self.lookback, 0])
            valid_indices.append(i + self.lookback - 1)
        return np.array(X), df, valid_indices

def fmt_pred(val):
    if val is None or (isinstance(val, float) and np.isnan(val)): return "N/A"
    return "涨" if int(val) == 1 else "跌"

def fmt_result(pred, actual):
    if pred is None or (isinstance(pred, float) and np.isnan(pred)): return "  "
    return "✓" if int(pred) == int(actual) else "✗"


from datetime import datetime
import json
import re
import time
import asyncio
from openai import AsyncOpenAI
import openai

# ==================== DeepSeek 异步并发核心逻辑 (双重锁绝对防 429) ====================

def get_log_time():
    return datetime.now().strftime("%H:%M:%S")

async def fetch_deepseek_single(client, messages, item_id, rate_lock, global_state, sem):
    """
    单次异步API调用：并发数控制 + 漏桶控速 + 指数退避重试
    """
    attempt = 0
    while True:
        attempt += 1
        
        # 0. 并发锁：保证同一时刻，最多只有 N 个请求在英伟达服务器上运行
        async with sem:
            # 1. 速率锁：保证就算允许并发，发包的瞬间也要间隔至少 2.1 秒 (约 28次/分钟，留足余量)
            async with rate_lock:
                now = time.time()
                elapsed = now - global_state['last_call_time']
                if elapsed < 2.1:
                    await asyncio.sleep(2.1 - elapsed)
                global_state['last_call_time'] = time.time()
                
            print(f"[{get_log_time()}] 🚀 [任务 {item_id}] 发起请求 (第 {attempt} 次尝试)...")
                
            try:
                # 2. 调用 API (注意：请确保 model 是你刚才探测到的最新可用模型名)
                completion = await client.chat.completions.create(
                    model="deepseek-ai/deepseek-v3.2", # <--- 如果旧名字报错410，请务必替换为最新名字
                    messages=messages,
                    temperature=0.6,
                    top_p=0.7,
                    max_tokens=1024,
                    timeout=60.0
                )
                
                content = completion.choices[0].message.content
                match = re.search(r'\{.*\}', content, re.DOTALL)
                
                if match:
                    res = json.loads(match.group(0))
                    pred = res.get('prediction')
                    
                    if pred in [0, 1]:
                        pred_str = "涨" if pred == 1 else "跌"
                        print(f"[{get_log_time()}] ✅ [任务 {item_id}] 成功获取预测: {pred_str}")
                        return item_id, int(pred)
                    else:
                        print(f"[{get_log_time()}] ⚠️ [任务 {item_id}] 异常预测值: {pred}，重新请求...")
                        continue
                else:
                    print(f"[{get_log_time()}] ⚠️ [任务 {item_id}] JSON 解析失败，重新请求...")
                    continue

            except openai.RateLimitError:
                # 指数退避：第一次等 5 秒，第二次等 10 秒，第三次等 20 秒...
                wait_time = 5 * (2 ** (attempt - 1))
                wait_time = min(wait_time, 60) # 最多等 60 秒
                print(f"[{get_log_time()}] 🛑 [任务 {item_id}] 触发限流 (429)，休眠 {wait_time} 秒后重试...")
                await asyncio.sleep(wait_time)
            except openai.APITimeoutError:
                print(f"[{get_log_time()}] ⏳ [任务 {item_id}] 请求超时 (>60s)，重新请求...")
                await asyncio.sleep(2)
            except openai.APIStatusError as e:
                if e.status_code in [400, 401, 403, 404, 410]:
                    print(f"[{get_log_time()}] 🚨 [任务 {item_id}] 致命错误 (HTTP {e.status_code}): {e.message}。放弃。")
                    return item_id, None
                print(f"[{get_log_time()}] ❌ [任务 {item_id}] API 报错 (HTTP {e.status_code})，休眠 3 秒后重试...")
                await asyncio.sleep(3)
            except Exception as e:
                print(f"[{get_log_time()}] 💥 [任务 {item_id}] 未知错误: {e}，休眠 3 秒后重新请求...")
                await asyncio.sleep(3)


async def run_all_deepseek_async(tasks_data, api_key):
    """统筹所有的异步请求"""
    total = len(tasks_data)
    
    # === 核心参数配置 ===
    MAX_CONCURRENT = 8 # 绝对并发上限：同一时刻最多 2 个请求在处理
    
    print("\n" + "="*80)
    print(f"[{get_log_time()}] 开始 DeepSeek 异步并发请求 (双重防 429 版)")
    print(f"总任务数: {total} 个 | 严格并发数: {MAX_CONCURRENT} | 发包间隔: 2.1 秒")
    print("="*80 + "\n")
    
    client = AsyncOpenAI(
        base_url="https://integrate.api.nvidia.com/v1",
        api_key=api_key
    )
    
    sem = asyncio.Semaphore(MAX_CONCURRENT)
    rate_lock = asyncio.Lock()
    global_state = {'last_call_time': 0.0}
    ds_results_dict = {}
    
    tasks = []
    for item in tasks_data:
        task = asyncio.create_task(
            fetch_deepseek_single(client, item['messages'], item['id'], rate_lock, global_state, sem)
        )
        tasks.append(task)
        
    completed = 0
    for coro in asyncio.as_completed(tasks):
        item_id, pred = await coro
        ds_results_dict[item_id] = pred
        completed += 1
        
        if completed % 10 == 0 or completed == total:
            print(f"\n[{get_log_time()}] 📊 【全局进度】 DeepSeek 已完成: {completed}/{total} ({(completed/total):.1%})\n")
            
    return ds_results_dict
# ==================== 主程序 ====================

def main():
    print("=" * 115)
    print("BTC 1h 涨跌预测准确率测试 (DLinear, N-BEATS, TTM-R2, DeepSeek-R1 异步版)")
    print("=" * 115)

    device = torch.device('mps' if torch.backends.mps.is_available() else 'cpu')
    print(f"使用设备: {device}")

    # ========== 步骤1: 加载模型 ==========
    print("\n[1/4] 加载本地模型...")
    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']).to(device).eval()
    dlinear_model.load_state_dict(dlinear_checkpoint['model_state_dict'])
    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']).to(device).eval()
    with open('./models/checkpoints/nbeats_btc_1h/preprocessor.pkl', 'rb') as f:
        nbeats_preprocessor = pickle.load(f)

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

    ds_api_key = os.environ.get("NVIDIA_API_KEY", "")
    if not ds_api_key:
        print("  ⚠️ 未检测到 NVIDIA_API_KEY，跳过 DeepSeek 测试。")
    else:
        print("  ✓ DeepSeek-R1 (异步并发) 已就绪")

    # ========== 步骤2: 获取历史数据 ==========
    print("\n[2/4] 获取历史数据并计算指标...")
    fetcher = BinanceDataFetcher()
    df = fetcher.fetch_ohlcv('BTC', '1h', start_date='2025-10-01', end_date='2026-03-01', use_cache=False)
    
    df['rsi'] = ta.rsi(df['close'], length=14)
    macd = ta.macd(df['close'])
    if macd is not None and not macd.empty:
        df['macd_hist'] = macd.iloc[:, 1]

    # ========== 步骤3: 本地预处理与预测 ==========
    print("\n[3/4] 运行本地预测 & 构建 API 任务队列...")
    # DLinear预处理
    df_dl = df.copy()
    df_dl['rsi'] = ta.rsi(df_dl['close'], length=14)
    macd_dl = ta.macd(df_dl['close'])
    if macd_dl is not None and not macd_dl.empty:
        df_dl['macd'], df_dl['macd_hist'], df_dl['macd_signal'] = macd_dl.iloc[:, 0], macd_dl.iloc[:, 1], macd_dl.iloc[:, 2]
    df_dl = df_dl.dropna().reset_index(drop=True)
    scaled_dl = dlinear_preprocessor.scaler.transform(df_dl[['open', 'high', 'low', 'close', 'volume', 'rsi', 'macd', 'macd_hist', 'macd_signal']])
    X_dl, dl_ts_map = [], {}
    for i in range(len(scaled_dl) - dlinear_preprocessor.lookback):
        X_dl.append(scaled_dl[i:i + dlinear_preprocessor.lookback])
        dl_ts_map[df_dl.iloc[i + dlinear_preprocessor.lookback - 1]['timestamp']] = i
    X_dl = np.array(X_dl)

    # N-BEATS预处理
    X_nb, df_nb, idx_nb = nbeats_preprocessor.preprocess(df, fit=False)

    # TTM预处理
    scaler_ttm = StandardScaler()
    scaled_close_ttm = scaler_ttm.fit_transform(df[['close']])
    X_ttm, idx_ttm_map = [], {}
    for i in range(len(scaled_close_ttm) - 512):
        X_ttm.append(scaled_close_ttm[i:i+512])
        idx_ttm_map[i + 511] = i
    X_ttm = np.array(X_ttm)

    target_targets = []
    for target_idx in range(1, len(df)):
        target_time_bj = utc_to_beijing(df.iloc[target_idx]['timestamp'])
        if target_time_bj.year == 2026 and target_time_bj.month == 2:  # 锁定 2026年 2月 数据
            target_targets.append(target_idx)

    temp_results = []
    deepseek_tasks_data = []

    for target_idx in target_targets:
        input_idx = target_idx - 1
        current_close, target_close = df.iloc[input_idx]['close'], df.iloc[target_idx]['close']
        actual_label = 1 if target_close > current_close else 0

        # DL
        dl_pred, input_ts = None, df.iloc[input_idx]['timestamp']
        if (dl_i := dl_ts_map.get(input_ts)) is not None and dl_i < len(X_dl):
            with torch.no_grad():
                logits = dlinear_model(torch.FloatTensor(X_dl[dl_i:dl_i+1]).to(device)).cpu().numpy()
                dl_pred = int((1 / (1 + np.exp(-logits))).item() > 0.5)

        # NB
        nb_pred = None
        if input_idx in idx_nb:
            nb_i = idx_nb.index(input_idx)
            with torch.no_grad():
                nb_out = nbeats_model(torch.FloatTensor(X_nb[nb_i:nb_i+1]).to(device)).cpu().numpy()
                nb_pred = 1 if (nb_out[0] if nb_out.ndim == 1 else nb_out[0, 0]) > X_nb[nb_i, -1] else 0

        # TTM
        ttm_pred = None
        if (ttm_i := idx_ttm_map.get(input_idx)) is not None:
            with torch.no_grad():
                ttm_pred_scaled = ttm_model(torch.FloatTensor(X_ttm[ttm_i:ttm_i+1]).to(device)).prediction_outputs[0, 0, 0].cpu().numpy()
                ttm_pred = 1 if ttm_pred_scaled > X_ttm[ttm_i, -1, 0] else 0

        # 构造 DeepSeek Payload
        if ds_api_key:
            recent_k = df.iloc[max(0, input_idx-14) : input_idx+1]
            if len(recent_k) >= 5:
                kline_text = "".join([f"时间: {row['timestamp']}, 开盘: {row['open']:.2f}, 收盘: {row['close']:.2f}, RSI: {row.get('rsi', 0):.2f}\n" for _, row in recent_k.iterrows()])
                sys_prompt = '你是一位顶级量化分析师。请分析BTC 1小时K线，预测下一小时涨跌。必须输出合法的JSON格式：{"prediction": 1(看涨)或0(看跌), "reason": "理由"}'
                messages = [
                    {"role": "system", "content": sys_prompt}, 
                    {"role": "user", "content": f"近15周期数据:\n{kline_text}\n输出JSON:"}
                ]
                deepseek_tasks_data.append({'id': target_idx, 'messages': messages})

        temp_results.append({
            'id': target_idx, 'time_utc': df.iloc[target_idx]['timestamp'], 'time_bj': utc_to_beijing(df.iloc[target_idx]['timestamp']),
            'current_close': current_close, 'target_close': target_close, 'actual': actual_label,
            'dlinear': dl_pred, 'nbeats': nb_pred, 'ttm': ttm_pred, 'deepseek': None
        })

    # ========== 异步执行 DeepSeek 并合入数据 ==========
    if deepseek_tasks_data:
        # 并发数设为 5，对 NVIDIA 免费/基础层级比较安全
        ds_predictions = asyncio.run(run_all_deepseek_async(deepseek_tasks_data, ds_api_key))
        for res in temp_results:
            if res['id'] in ds_predictions:
                res['deepseek'] = ds_predictions[res['id']]

    # 结算集成模型结果 (四投票)
    final_results = []
    for res in temp_results:
        votes = [v for v in [res['dlinear'], res['nbeats'], res['ttm'], res['deepseek']] if v is not None]
        if votes:
            vote_sum = sum(votes)
            if vote_sum > len(votes) / 2: res['ensemble'] = 1
            elif vote_sum < len(votes) / 2: res['ensemble'] = 0
            else: res['ensemble'] = res['dlinear'] if res['dlinear'] is not None else 0
        else:
            res['ensemble'] = None
        final_results.append(res)

    # ========== 步骤4: 输出结果 ==========
    if not final_results:
        print("没有找到符合条件的测试样本")
        return

    df_res = pd.DataFrame(final_results)
    df_res['date_bj'] = df_res['time_bj'].dt.date
    df_res['us_open'] = df_res['time_utc'].apply(is_us_stock_open_by_utc)

    # === (下方打印逻辑与原版一致，缩略优化版) ===
    print(f"\n[4/4] 准确率统计 (北京时间)\n总样本数: {len(df_res)}")
    print("\n" + "=" * 115 + "\n每小时详细数据\n" + "=" * 115)
    hourly_h = f"{'北京时间':<18} {'前收盘':>10} {'现收盘':>10} {'实际':>4} {'DL':>3} {'':>1} {'NB':>3} {'':>1} {'TTM':>3} {'':>1} {'DS':>3} {'':>1} {'集成':>4} {'':<1}"

    for date_bj, day_group in df_res.groupby('date_bj'):
        us_str = "美股开盘" if day_group['us_open'].iloc[0] else "美股休市"
        def s_acc(c): return (day_group.dropna(subset=[c])[c] == day_group.dropna(subset=[c])['actual']).mean() if len(day_group.dropna(subset=[c])) > 0 else None
        print(f"\n--- {date_bj} ({us_str}) --- DL:{s_acc('dlinear'):.1%} | NB:{s_acc('nbeats'):.1%} | TTM:{s_acc('ttm'):.1%} | DS:{s_acc('deepseek'):.1%} | 集成:{s_acc('ensemble'):.1%}")
        print(hourly_h + "\n" + "-" * 115)
        for _, r in day_group.iterrows():
            print(f"{r['time_bj'].strftime('%m-%d %H:%M'):<18} {r['current_close']:>10.2f} {r['target_close']:>10.2f} {fmt_pred(r['actual']):>4} "
                  f"{fmt_pred(r['dlinear']):>3} {fmt_result(r['dlinear'],r['actual']):>1} {fmt_pred(r['nbeats']):>3} {fmt_result(r['nbeats'],r['actual']):>1} "
                  f"{fmt_pred(r['ttm']):>3} {fmt_result(r['ttm'],r['actual']):>1} {fmt_pred(r['deepseek']):>3} {fmt_result(r['deepseek'],r['actual']):>1} "
                  f"{fmt_pred(r['ensemble']):>4} {fmt_result(r['ensemble'],r['actual']):<1}")

    print("\n\n" + "=" * 115 + "\n每日准确率汇总 (北京时间)\n" + "=" * 115)
    print(f"{'日期':<14} {'样本':>4} {'美股':>6}  {'DLinear':>8} {'N-BEATS':>8} {'TTM-R2':>8} {'DeepSeek':>8} {'集成':>10}\n" + "-" * 115)
    for date_bj, g in df_res.groupby('date_bj'):
        def g_acc(c): return (g.dropna(subset=[c])[c] == g.dropna(subset=[c])['actual']).mean() if len(g.dropna(subset=[c])) > 0 else None
        print(f"{str(date_bj):<14} {len(g):>4} {'开盘' if g['us_open'].iloc[0] else '休市':>6}  {g_acc('dlinear'):>8.1%} {g_acc('nbeats'):>8.1%} {g_acc('ttm'):>8.1%} {g_acc('deepseek'):>8.1%} {g_acc('ensemble'):>10.1%}".replace('nan%', 'N/A'))

    def t_acc(c): return (df_res.dropna(subset=[c])[c] == df_res.dropna(subset=[c])['actual']).mean() if len(df_res.dropna(subset=[c])) > 0 else None
    print("-" * 115)
    print(f"{'总计':<14} {len(df_res):>4} {'':>6}  {t_acc('dlinear'):>8.2%} {t_acc('nbeats'):>8.2%} {t_acc('ttm'):>8.2%} {t_acc('deepseek'):>8.2%} {t_acc('ensemble'):>10.2%}".replace('nan%', 'N/A'))
    print("=" * 115)

if __name__ == "__main__":
    main()