"""加密货币价格预测系统 - DLinear模型端到端测试"""

import sys
import os

# 添加项目根目录到路径
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from data.fetcher import BinanceDataFetcher
from data.preprocessor import CryptoPreprocessor, split_data
from data.dataset import create_dataloaders
from models.dlinear.model import DLinear, DLinearConfig
from training.trainer import Trainer
from training.config import TrainingConfig

def main():
    print("="*80)
    print("加密货币价格预测系统 - DLINEAR模型测试")
    print("="*80)

    # 步骤1: 获取数据
    print("\n[1/7] 从Binance获取BTC 1小时数据...")
    fetcher = BinanceDataFetcher()
    df = fetcher.fetch_ohlcv('BTC', '1h', use_cache=True)
    print(f"✓ 已获取 {len(df)} 行数据，时间范围：{df['timestamp'].min()} 至 {df['timestamp'].max()}")

    # 步骤2: 数据预处理
    print("\n[2/7] 数据预处理（添加技术指标、创建序列）...")
    preprocessor = CryptoPreprocessor(lookback=168, forecast_horizon=24)
    X, y, processed_df = preprocessor.preprocess_pipeline(df, task='classification', fit=True)
    print(f"✓ 已创建序列：X形状 = {X.shape}, y形状 = {y.shape}")
    print(f"✓ 特征数量：{X.shape[-1]}")
    print(f"✓ 标签分布：下跌(0) = {(y==0).sum()}, 上涨(1) = {(y==1).sum()}")

    # 步骤3: 划分数据集
    print("\n[3/7] 划分训练/验证/测试集 (70/15/15)...")
    train_data, val_data, test_data = split_data(X, y, train_split=0.7, val_split=0.15)

    # 步骤4: 创建数据加载器
    print("\n[4/7] 创建PyTorch数据加载器...")
    train_loader, val_loader, test_loader = create_dataloaders(
        train_data, val_data, test_data,
        task='classification',
        batch_size=64
    )
    print(f"✓ 训练批次：{len(train_loader)}")
    print(f"✓ 验证批次：{len(val_loader)}")
    print(f"✓ 测试批次：{len(test_loader)}")

    # 步骤5: 创建模型
    print("\n[5/7] 创建DLinear模型...")
    model_config = DLinearConfig(
        input_dim=X.shape[-1],
        lookback=168,
        task='classification',
        hidden_dim=512,
        dropout=0.1,
        moving_avg_kernel=25
    )
    model = DLinear(model_config)
    print(f"✓ 模型已创建，可训练参数：{model.count_parameters():,} 个")

    # 步骤6: 训练模型
    print("\n[6/7] 开始训练模型...")
    train_config = TrainingConfig(
        batch_size=64,
        max_epochs=20,  # 测试用较少轮数
        learning_rate=0.001,
        early_stopping_patience=5,
        lr_scheduler='plateau',
        save_best=True,
        checkpoint_dir='./models/checkpoints/dlinear_test'
    )

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

    # 步骤7: 评估模型
    print("\n[7/7] 在测试集上评估模型...")
    metrics = trainer.evaluate(test_loader)

    # 打印结果摘要
    print("\n" + "="*80)
    print("最终测试结果")
    print("="*80)
    print(f"\n分类指标:")
    print(f"  准确率:  {metrics.get('accuracy', 0):.4f}")
    print(f"  F1分数:  {metrics.get('f1', 0):.4f}")
    print(f"  精确率:  {metrics.get('precision', 0):.4f}")
    print(f"  召回率:  {metrics.get('recall', 0):.4f}")

    if 'true_positives' in metrics:
        print(f"\n混淆矩阵:")
        print(f"  真阳性:  {metrics['true_positives']}")
        print(f"  真阴性:  {metrics['true_negatives']}")
        print(f"  假阳性:  {metrics['false_positives']}")
        print(f"  假阴性:  {metrics['false_negatives']}")

    print(f"\n回测结果:")
    print(f"  总收益率:    {metrics.get('backtest_total_return', 0)*100:.2f}%")
    print(f"  最终资金:    ${metrics.get('backtest_final_capital', 0):,.2f}")
    print(f"  夏普比率:    {metrics.get('backtest_sharpe_ratio', 0):.4f}")
    print(f"  最大回撤:    {metrics.get('backtest_max_drawdown', 0)*100:.2f}%")
    print(f"  胜率:        {metrics.get('backtest_win_rate', 0)*100:.2f}%")
    print(f"  交易次数:    {metrics.get('backtest_num_trades', 0)}")

    # 测试ONNX导出
    print("\n" + "="*80)
    print("测试ONNX导出")
    print("="*80)

    try:
        import onnx
        import onnxruntime as ort
        import torch
        import numpy as np

        onnx_path = './models/checkpoints/dlinear_test/model.onnx'
        os.makedirs(os.path.dirname(onnx_path), exist_ok=True)

        print("\n[1/3] 导出模型到ONNX格式...")
        model.to_onnx(onnx_path)
        print(f"✓ ONNX模型已保存至：{onnx_path}")

        print("\n[2/3] 验证ONNX模型...")
        onnx_model = onnx.load(onnx_path)
        onnx.checker.check_model(onnx_model)
        print("✓ ONNX模型有效")

        print("\n[3/3] 测试ONNX推理与PyTorch对比...")
        # 创建测试输入
        test_input = torch.randn(1, 168, X.shape[-1])

        # PyTorch推理（在CPU上）
        model.cpu()
        model.eval()
        with torch.no_grad():
            pytorch_output = model(test_input).numpy()

        # ONNX推理
        ort_session = ort.InferenceSession(onnx_path)
        onnx_output = ort_session.run(
            None,
            {'input': test_input.numpy()}
        )[0]

        # 比较输出
        diff = np.abs(pytorch_output - onnx_output).max()
        print(f"✓ PyTorch与ONNX的最大差异：{diff:.8f}")

        if diff < 1e-5:
            print("✓ ONNX导出成功！输出在容差范围内匹配。")
        else:
            print("⚠ 警告：ONNX输出与PyTorch有差异（但可能仍可接受）")

    except Exception as e:
        print(f"✗ ONNX导出失败：{e}")

    print("\n" + "="*80)
    print("测试成功完成！")
    print("="*80)
    print("\n下一步:")
    print("  1. 查看训练历史：./models/checkpoints/dlinear_test/")
    print("  2. 实现N-BEATS模型（阶段3）")
    print("  3. 在所有币种上训练（BTC、ETH、SOL）")
    print("  4. 创建FastAPI预测服务")

if __name__ == "__main__":
    main()