"""直接预测未来1小时的价格（而非收益率）"""

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.dlinear.model import DLinear, DLinearConfig
from training.trainer import Trainer
from training.config import TrainingConfig
import json
import pickle

class PricePreprocessor:
    """预测未来价格的预处理器"""

    def __init__(self, lookback=168, forecast_horizon=1):
        self.lookback = lookback
        self.forecast_horizon = forecast_horizon
        self.scaler = StandardScaler()
        self.price_scaler = StandardScaler()  # 单独的价格缩放器
        self.is_fitted = False

    def add_features(self, df):
        """添加技术指标"""
        df = df.copy()

        # RSI
        df['rsi'] = ta.rsi(df['close'], length=14)

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

        return df

    def create_labels(self, df):
        """创建标签：未来1小时的价格（绝对价格）"""
        df = df.copy()
        # 标签是未来1小时的收盘价
        df['future_price'] = df['close'].shift(-self.forecast_horizon)
        df['label'] = df['future_price']
        return df[:-self.forecast_horizon]

    def preprocess(self, df, fit=True):
        """完整预处理流程"""
        # 添加技术指标
        df = self.add_features(df)
        df = self.create_labels(df)

        # 删除NaN
        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])
            # 单独标准化价格标签（用于训练）
            scaled_labels = self.price_scaler.fit_transform(df[['label']])
            self.is_fitted = True
            self.feature_columns = feature_cols
        else:
            if not self.is_fitted:
                raise ValueError("Must fit before transform")
            scaled_features = self.scaler.transform(df[feature_cols])
            scaled_labels = self.price_scaler.transform(df[['label']])

        # 创建序列
        X, y = [], []
        for i in range(len(scaled_features) - self.lookback):
            X.append(scaled_features[i:i + self.lookback])
            y.append(scaled_labels[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("训练模型：直接预测未来1小时价格")
    print("="*80)

    # ========== 步骤1: 获取数据 ==========
    print("\n[1/7] 获取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/7] 数据预处理（预测绝对价格）...")
    preprocessor = PricePreprocessor(lookback=168, forecast_horizon=1)
    X, y, processed_df = preprocessor.preprocess(df, fit=True)

    print(f"✓ 序列数量: {len(X):,}")
    print(f"✓ 序列形状: {X.shape}")
    print(f"✓ 特征数量: {len(preprocessor.feature_columns)}")

    # 还原真实价格看看统计
    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: 划分数据集（95/5）==========
    print("\n[3/7] 划分数据集（95%训练 + 5%验证）...")
    n = len(X)
    train_end = int(n * 0.95)

    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/7] 创建数据加载器...")
    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=64, shuffle=False)
    val_loader = DataLoader(val_dataset, batch_size=64, shuffle=False)

    # ========== 步骤5: 创建模型 ==========
    print("\n[5/7] 创建DLinear模型...")
    model_config = DLinearConfig(
        input_dim=9,
        lookback=168,
        task='regression',  # 仍然是回归，但预测价格而非收益率
        hidden_dim=256,
        dropout=0.1,
        moving_avg_kernel=25
    )
    model = DLinear(model_config)
    print(f"✓ 模型参数: {model.count_parameters():,}")
    print(f"✓ 预测目标: 未来1小时的绝对价格")

    # ========== 步骤6: 训练 ==========
    print("\n[6/7] 开始训练...")
    train_config = TrainingConfig(
        batch_size=64,
        max_epochs=50,
        learning_rate=0.001,
        early_stopping_patience=15,
        lr_scheduler='plateau',
        save_best=True,
        checkpoint_dir='./models/checkpoints/dlinear_predict_price'
    )

    trainer = Trainer(model, train_config, task='regression')
    history = trainer.fit(train_loader, val_loader)

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

    # 在验证集上预测
    device = torch.device('mps' if torch.backends.mps.is_available() else 'cpu')
    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).squeeze()
            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  # +168是lookback
    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("训练完成 - 验证集性能")
    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}%")
    print(f"  预测上涨: {(pred_direction==1).sum()} / 实际上涨: {(true_direction==1).sum()}")
    print(f"  预测下跌: {(pred_direction==0).sum()} / 实际下跌: {(true_direction==0).sum()}")

    print(f"\n💡 价格预测示例（前10个）:")
    print(f"{'当前价格':<12} {'预测价格':<12} {'实际价格':<12} {'预测方向':<8} {'实际方向':<8}")
    print("-"*70)
    for i in range(min(10, len(pred_prices))):
        pred_dir = "上涨" if pred_prices[i] > current_prices[i] else "下跌"
        true_dir = "上涨" if true_prices[i] > current_prices[i] else "下跌"
        print(f"${current_prices[i]:<11,.2f} ${pred_prices[i]:<11,.2f} ${true_prices[i]:<11,.2f} "
              f"{pred_dir:<8} {true_dir:<8}")

    # 保存模型和预处理器
    with open('./models/checkpoints/dlinear_predict_price/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)
    }
    with open('./models/checkpoints/dlinear_predict_price/metrics.json', 'w') as f:
        json.dump(metrics, f, indent=2)

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

if __name__ == "__main__":
    main()