"""使用N-BEATS模型训练：预测未来1小时BTC价格"""

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

import numpy as np
import pandas as pd
import pandas_ta as ta
from sklearn.preprocessing import StandardScaler
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset

from data.fetcher import BinanceDataFetcher
from models.nbeats.model import NBeats, NBeatsConfig
import json
import pickle


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:
            if not self.is_fitted:
                raise ValueError("Must fit before transform")
            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 train_nbeats():
    print("="*80)
    print("N-BEATS模型训练 - 预测未来1小时BTC价格")
    print("="*80)

    # ========== 步骤1: 获取数据 ==========
    print("\n[1/8] 获取BTC 2020-2025年数据...")
    fetcher = BinanceDataFetcher()
    df = fetcher.fetch_ohlcv(
        symbol='BTC',
        timeframe='1h',
        start_date='2020-01-01',
        end_date='2025-12-31',
        use_cache=True
    )

    print(f"✓ 数据量: {len(df):,} 行")
    print(f"✓ 时间: {df['timestamp'].min()} 至 {df['timestamp'].max()}")
    print(f"✓ 价格: ${df['close'].iloc[0]:,.2f} → ${df['close'].iloc[-1]:,.2f}")

    # ========== 步骤2: 预处理 ==========
    print("\n[2/8] 数据预处理...")
    preprocessor = NBeatsPreprocessor(lookback=168, forecast_horizon=1)
    X, y, processed_df = preprocessor.preprocess(df, fit=True)

    print(f"✓ 序列数量: {len(X):,}")
    print(f"✓ 输入形状: {X.shape}")
    print(f"✓ 输出形状: {y.shape}")

    # 还原真实价格看统计
    y_real = preprocessor.inverse_transform_price(y)
    print(f"\n✓ 目标价格统计:")
    print(f"   最小值: ${y_real.min():,.2f}")
    print(f"   最大值: ${y_real.max():,.2f}")
    print(f"   均值:   ${y_real.mean():,.2f}")

    # ========== 步骤3: 划分数据集 ==========
    print("\n[3/8] 划分数据集（90%训练 + 10%验证）...")
    n = len(X)
    train_end = int(n * 0.9)

    X_train, y_train = X[:train_end], y[:train_end]
    X_val, y_val = X[train_end:], y[train_end:]

    print(f"  训练集: {len(X_train):,} 样本")
    print(f"  验证集: {len(X_val):,} 样本")

    # ========== 步骤4: 创建数据加载器 ==========
    print("\n[4/8] 创建数据加载器...")
    train_dataset = TensorDataset(
        torch.FloatTensor(X_train),
        torch.FloatTensor(y_train)
    )
    val_dataset = TensorDataset(
        torch.FloatTensor(X_val),
        torch.FloatTensor(y_val)
    )

    train_loader = DataLoader(train_dataset, batch_size=128, shuffle=True)
    val_loader = DataLoader(val_dataset, batch_size=128, shuffle=False)

    # ========== 步骤5: 创建N-BEATS模型 ==========
    print("\n[5/8] 创建N-BEATS模型...")
    config = NBeatsConfig(
        input_size=168,
        forecast_size=1,
        stacks=['trend', 'seasonality', 'generic'],
        num_blocks_per_stack=3,
        block_layers=4,
        layer_size=256,
        trend_degree=2,
        seasonality_harmonics=4
    )
    model = NBeats(config)

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

    print(f"✓ 模型参数: {model.count_parameters():,}")
    print(f"✓ 设备: {device}")
    print(f"✓ 模型架构: Trend + Seasonality + Generic")

    # ========== 步骤6: 训练 ==========
    print("\n[6/8] 开始训练...")

    criterion = nn.MSELoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
        optimizer, mode='min', factor=0.5, patience=5, verbose=True
    )

    best_val_loss = float('inf')
    patience = 15
    patience_counter = 0
    max_epochs = 50

    for epoch in range(max_epochs):
        # 训练
        model.train()
        train_loss = 0
        for batch_X, batch_y in train_loader:
            batch_X, batch_y = batch_X.to(device), batch_y.to(device)

            optimizer.zero_grad()
            outputs = model(batch_X)
            loss = criterion(outputs, batch_y)
            loss.backward()
            optimizer.step()

            train_loss += loss.item()

        train_loss /= len(train_loader)

        # 验证
        model.eval()
        val_loss = 0
        with torch.no_grad():
            for batch_X, batch_y in val_loader:
                batch_X, batch_y = batch_X.to(device), batch_y.to(device)
                outputs = model(batch_X)
                loss = criterion(outputs, batch_y)
                val_loss += loss.item()

        val_loss /= len(val_loader)

        # 学习率调度
        scheduler.step(val_loss)

        print(f"Epoch {epoch+1}/{max_epochs}")
        print(f"  Train loss: {train_loss:.6f}")
        print(f"  Val loss: {val_loss:.6f}")
        print(f"  LR: {optimizer.param_groups[0]['lr']:.6f}")

        # 保存最佳模型
        if val_loss < best_val_loss:
            best_val_loss = val_loss
            patience_counter = 0
            os.makedirs('./models/checkpoints/nbeats_btc_1h', exist_ok=True)
            model.save('./models/checkpoints/nbeats_btc_1h/best_model.pt')
            print(f"  Saved best model (val_loss: {val_loss:.6f})")
        else:
            patience_counter += 1

        # Early stopping
        if patience_counter >= patience:
            print(f"\nEarly stopping at epoch {epoch+1}")
            break

    # ========== 步骤7: 评估 ==========
    print("\n[7/8] 评估模型...")

    # 重新加载最佳模型
    checkpoint = torch.load('./models/checkpoints/nbeats_btc_1h/best_model.pt', map_location=device)
    model = NBeats(checkpoint['config'])
    model.load_state_dict(checkpoint['model_state_dict'])
    model = model.to(device)
    model.eval()

    val_preds = []
    with torch.no_grad():
        for batch_X, _ in val_loader:
            batch_X = batch_X.to(device)
            pred = model(batch_X)
            val_preds.extend(pred.cpu().numpy())

    val_preds = np.array(val_preds)

    # 还原真实价格
    pred_prices = preprocessor.inverse_transform_price(val_preds)
    true_prices = preprocessor.inverse_transform_price(y_val)

    # 当前价格
    val_start_idx = train_end + 168
    current_prices = processed_df['close'].iloc[val_start_idx:val_start_idx+len(pred_prices)].values

    # 方向准确率
    pred_direction = (pred_prices > current_prices).astype(int)
    true_direction = (true_prices > current_prices).astype(int)
    direction_acc = (pred_direction == true_direction).mean()

    # 价格预测误差
    price_mae = np.abs(pred_prices - true_prices).mean()
    price_rmse = np.sqrt(((pred_prices - true_prices) ** 2).mean())
    price_mape = np.mean(np.abs((true_prices - pred_prices) / true_prices)) * 100

    print("\n" + "="*80)
    print("训练完成 - N-BEATS模型性能")
    print("="*80)

    print(f"\n📊 价格预测误差:")
    print(f"  MAE:  ${price_mae:,.2f}")
    print(f"  RMSE: ${price_rmse:,.2f}")
    print(f"  MAPE: {price_mape:.2f}%")

    print(f"\n📈 方向预测:")
    print(f"  方向准确率: {direction_acc*100:.2f}%")

    # ========== 步骤8: 保存 ==========
    print("\n[8/8] 保存模型和预处理器...")

    with open('./models/checkpoints/nbeats_btc_1h/preprocessor.pkl', 'wb') as f:
        pickle.dump(preprocessor, f)

    metrics = {
        'price_mae': float(price_mae),
        'price_rmse': float(price_rmse),
        'price_mape': float(price_mape),
        'direction_accuracy': float(direction_acc),
        'best_val_loss': float(best_val_loss)
    }
    with open('./models/checkpoints/nbeats_btc_1h/metrics.json', 'w') as f:
        json.dump(metrics, f, indent=2)

    print(f"\n✓ 模型保存至: ./models/checkpoints/nbeats_btc_1h/")
    print("="*80)


if __name__ == "__main__":
    train_nbeats()