"""训练DLinear模型 - 预测未来1小时涨跌（已修复双重Sigmoid Bug）"""

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
import json
import numpy as np

def main():
    print("="*80)
    print("DLinear模型训练 - 预测未来1小时涨跌（已修复双重Sigmoid Bug）")
    print("="*80)

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

    # 步骤2: 数据预处理（预测未来1小时）
    print("\n[2/7] 数据预处理（预测未来1小时的涨跌）...")
    preprocessor = CryptoPreprocessor(
        lookback=168,           # 输入：过去168小时（7天）
        forecast_horizon=1,     # ← 改为预测未来1小时
        price_change_threshold=0.0  # 任何涨幅>0都算上涨
    )
    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]}")

    # 计算类别分布
    n_down = (y == 0).sum()
    n_up = (y == 1).sum()
    print(f"✓ 标签分布：下跌(0) = {n_down} ({n_down/len(y)*100:.1f}%), 上涨(1) = {n_up} ({n_up/len(y)*100:.1f}%)")

    # 计算类别权重
    pos_weight = n_down / n_up if n_up > 0 else 1.0
    print(f"✓ 计算类别权重：pos_weight = {pos_weight:.2f}")

    # 步骤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)

    print(f"  训练集：下跌 {(train_data[1]==0).sum()}, 上涨 {(train_data[1]==1).sum()}")
    print(f"  验证集：下跌 {(val_data[1]==0).sum()}, 上涨 {(val_data[1]==1).sum()}")
    print(f"  测试集：下跌 {(test_data[1]==0).sum()}, 上涨 {(test_data[1]==1).sum()}")

    # 步骤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
    )

    # 步骤5: 创建模型
    print("\n[5/7] 创建DLinear模型（已修复双重Sigmoid Bug）...")
    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():,} 个")
    print(f"✓ Bug修复：输出层使用Identity（输出logits），由BCEWithLogitsLoss处理Sigmoid")

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

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

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

    # 保存评估结果
    results_path = './models/checkpoints/dlinear_btc_1h_forecast/test_metrics.json'
    with open(results_path, 'w') as f:
        json.dump(metrics, f, indent=2)

    print(f"\n✓ 评估结果已保存至：{results_path}")

    # 打印结果摘要
    print("\n" + "="*80)
    print("训练完成 - 最终测试结果（预测未来1小时涨跌）")
    print("="*80)

    print(f"\n📊 分类指标:")
    print(f"  准确率:  {metrics.get('accuracy', 0):.4f} ({metrics.get('accuracy', 0)*100:.2f}%)")
    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)}")

    # 性能分析
    print(f"\n🎯 性能评估:")
    acc = metrics.get('accuracy', 0)
    if acc > 0.55:
        print(f"  ✅ 准确率 {acc*100:.1f}% > 55%，模型有一定预测能力")
    elif acc > 0.50:
        print(f"  ⚠️  准确率 {acc*100:.1f}% 略优于随机（50%），预测能力较弱")
    else:
        print(f"  ❌ 准确率 {acc*100:.1f}% ≤ 50%，模型预测能力不足")

    total_return = metrics.get('backtest_total_return', 0)
    if total_return > 0:
        print(f"  ✅ 回测收益 {total_return*100:.2f}% > 0，策略有盈利潜力")
    else:
        print(f"  ❌ 回测收益 {total_return*100:.2f}% ≤ 0，策略亏损")

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

if __name__ == "__main__":
    main()