"""8:1: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()
        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):
        """创建标签：未来1小时的价格"""
        df = df.copy()
        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)
        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("8:1:1数据集划分 - 直接预测未来1小时价格")
    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] 数据预处理（只使用OHLCV + RSI + MACD）...")
    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"✓ 特征列表: {preprocessor.feature_columns}")

    # ========== 步骤3: 8:1:1划分数据集 ==========
    print("\n[3/8] 8:1:1划分数据集...")
    n = len(X)
    train_end = int(n * 0.8)
    val_end = int(n * 0.9)

    X_train = X[:train_end]
    y_train = y[:train_end]
    X_val = X[train_end:val_end]
    y_val = y[train_end:val_end]
    X_test = X[val_end:]
    y_test = y[val_end:]

    print(f"  总样本数: {n:,}")
    print(f"  训练集: {len(X_train):,} ({len(X_train)/n*100:.1f}%)")
    print(f"  验证集: {len(X_val):,} ({len(X_val)/n*100:.1f}%)")
    print(f"  测试集: {len(X_test):,} ({len(X_test)/n*100:.1f}%)")

    # 分析各数据集时间范围
    lookback = 168
    train_df = processed_df.iloc[lookback:train_end+lookback]
    val_df = processed_df.iloc[train_end+lookback:val_end+lookback]
    test_df = processed_df.iloc[val_end+lookback:]

    print(f"\n  时间范围分析:")
    print(f"  训练集: {train_df['timestamp'].min()} 至 {train_df['timestamp'].max()}")
    train_return = (train_df['close'].iloc[-1] - train_df['close'].iloc[0]) / train_df['close'].iloc[0]
    print(f"          ${train_df['close'].iloc[0]:,.0f} → ${train_df['close'].iloc[-1]:,.0f} ({train_return*100:+.1f}%)")

    print(f"  验证集: {val_df['timestamp'].min()} 至 {val_df['timestamp'].max()}")
    val_return = (val_df['close'].iloc[-1] - val_df['close'].iloc[0]) / val_df['close'].iloc[0]
    print(f"          ${val_df['close'].iloc[0]:,.0f} → ${val_df['close'].iloc[-1]:,.0f} ({val_return*100:+.1f}%)")

    print(f"  测试集: {test_df['timestamp'].min()} 至 {test_df['timestamp'].max()}")
    test_return = (test_df['close'].iloc[-1] - test_df['close'].iloc[0]) / test_df['close'].iloc[0]
    print(f"          ${test_df['close'].iloc[0]:,.0f} → ${test_df['close'].iloc[-1]:,.0f} ({test_return*100:+.1f}%)")

    # ========== 步骤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))
    test_dataset = TensorDataset(torch.FloatTensor(X_test), torch.FloatTensor(y_test))

    train_loader = DataLoader(train_dataset, batch_size=64, shuffle=False)
    val_loader = DataLoader(val_dataset, batch_size=64, shuffle=False)
    test_loader = DataLoader(test_dataset, batch_size=64, shuffle=False)

    # ========== 步骤5: 创建模型 ==========
    print("\n[5/8] 创建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():,}")

    # ========== 步骤6: 训练 ==========
    print("\n[6/8] 开始训练...")
    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_811_split'
    )

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

    # ========== 步骤7: 在测试集上评估 ==========
    print("\n[7/8] 在测试集上评估...")

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

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

    test_preds = np.array(test_preds)

    # 还原为真实价格
    pred_prices = preprocessor.inverse_transform_price(test_preds)
    true_prices = preprocessor.inverse_transform_price(y_test)

    # 获取当前价格
    test_start_idx = val_end + lookback
    current_prices = processed_df['close'].iloc[test_start_idx:test_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

    from sklearn.metrics import confusion_matrix
    cm = confusion_matrix(true_direction, pred_direction)
    tn, fp, fn, tp = cm.ravel()

    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详细分析:")
    print(f"  ✓ 预测涨且真涨: {tp} ({tp/len(pred_prices)*100:.1f}%)")
    print(f"  ✓ 预测跌且真跌: {tn} ({tn/len(pred_prices)*100:.1f}%)")
    print(f"  ✗ 预测涨但真跌: {fp} ({fp/len(pred_prices)*100:.1f}%)")
    print(f"  ✗ 预测跌但真涨: {fn} ({fn/len(pred_prices)*100:.1f}%)")

    # 性能评估
    print("\n" + "="*80)
    print("🎯 性能评估")
    print("="*80)

    if direction_acc > 0.55:
        print(f"\n  ✅ 方向准确率 {direction_acc*100:.2f}% > 55%，模型有预测能力")
    elif direction_acc > 0.52:
        print(f"\n  ⚠️  方向准确率 {direction_acc*100:.2f}% > 52%，略优于随机")
    elif direction_acc >= 0.48:
        print(f"\n  ⚠️  方向准确率 {direction_acc*100:.2f}% ≈ 50%，接近随机")
    else:
        print(f"\n  ❌ 方向准确率 {direction_acc*100:.2f}% < 48%，低于随机")

    # ========== 步骤8: 保存 ==========
    print("\n[8/8] 保存模型和结果...")

    with open('./models/checkpoints/dlinear_811_split/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),
        'true_positives': int(tp),
        'true_negatives': int(tn),
        'false_positives': int(fp),
        'false_negatives': int(fn)
    }
    with open('./models/checkpoints/dlinear_811_split/test_metrics.json', 'w') as f:
        json.dump(metrics, f, indent=2)

    print(f"\n✓ 模型保存至: ./models/checkpoints/dlinear_811_split/")

    # 对比分析
    print("\n" + "="*80)
    print("📊 数据集划分对比")
    print("="*80)

    print("\n之前（95:5划分，无独立测试集）:")
    print("  训练集: 49,733 (95%)")
    print("  验证集: 2,618 (5%)")
    print("  测试集: 无")
    print("  问题: 过拟合，验证集太小")

    print(f"\n当前（8:1:1划分，有独立测试集）:")
    print(f"  训练集: {len(X_train):,} (80%)")
    print(f"  验证集: {len(X_val):,} (10%)")
    print(f"  测试集: {len(X_test):,} (10%)")
    print(f"  优势: 有足够大的独立测试集评估性能")

    print(f"\n测试集性能:")
    print(f"  方向准确率: {direction_acc*100:.2f}%")
    print(f"  价格误差MAPE: {price_mape:.2f}%")

    print("\n" + "="*80)
    print("✓ 训练完成！")
    print("="*80)

if __name__ == "__main__":
    main()