"""简化特征训练：只使用OHLCV + RSI + MACD"""

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 SimplePreprocessor:
    """简化的预处理器：只使用OHLCV + RSI + MACD"""

    def __init__(self, lookback=168, forecast_horizon=1):
        self.lookback = lookback
        self.forecast_horizon = forecast_horizon
        self.scaler = StandardScaler()
        self.is_fitted = False

    def add_features(self, df):
        """添加RSI和MACD特征"""
        df = df.copy()

        # RSI (14 periods)
        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()
        df['future_price'] = df['close'].shift(-self.forecast_horizon)
        df['future_return'] = (df['future_price'] - df['close']) / df['close']
        df['label'] = df['future_return']
        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)

        # 选择特征列：OHLCV + RSI + MACD
        feature_cols = ['open', 'high', 'low', 'close', 'volume',
                       'rsi', 'macd', 'macd_hist', 'macd_signal']

        # 标准化
        if fit:
            scaled = self.scaler.fit_transform(df[feature_cols])
            self.is_fitted = True
            self.feature_columns = feature_cols
        else:
            if not self.is_fitted:
                raise ValueError("Must fit before transform")
            scaled = self.scaler.transform(df[feature_cols])

        # 创建序列
        X, y = [], []
        for i in range(len(scaled) - self.lookback):
            X.append(scaled[i:i + self.lookback])
            y.append(df['label'].iloc[i + self.lookback])

        return np.array(X), np.array(y), df

def main():
    print("="*80)
    print("简化特征训练：OHLCV + RSI + MACD")
    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] 数据预处理（只使用9个特征）...")
    preprocessor = SimplePreprocessor(lookback=168, forecast_horizon=1)
    X, y, processed_df = preprocessor.preprocess(df, fit=True)

    print(f"✓ 序列数量: {len(X):,}")
    print(f"✓ 序列形状: {X.shape}")
    print(f"✓ 特征列表:")
    for i, col in enumerate(preprocessor.feature_columns, 1):
        print(f"   {i}. {col}")

    print(f"\n✓ 标签统计:")
    print(f"   收益率范围: {y.min()*100:.4f}% ~ {y.max()*100:.4f}%")
    print(f"   均值: {y.mean()*100:.4f}%")
    print(f"   上涨样本: {(y>0).sum()} ({(y>0).sum()/len(y)*100:.1f}%)")
    print(f"   下跌样本: {(y<=0).sum()} ({(y<=0).sum()/len(y)*100:.1f}%)")

    # ========== 步骤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,  # OHLCV + RSI + MACD(3个)
        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/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_simple_features'
    )

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

    # ========== 步骤7: 评估并保存 ==========
    print("\n[7/7] 评估并保存...")
    val_metrics = trainer.evaluate(val_loader)

    # 保存结果
    with open('./models/checkpoints/dlinear_simple_features/metrics.json', 'w') as f:
        json.dump(val_metrics, f, indent=2)

    # 保存预处理器
    with open('./models/checkpoints/dlinear_simple_features/preprocessor.pkl', 'wb') as f:
        pickle.dump(preprocessor, f)

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

    print(f"\n📊 验证集性能:")
    print(f"  RMSE:         {val_metrics.get('rmse', 0)*100:.4f}%")
    print(f"  MAE:          {val_metrics.get('mae', 0)*100:.4f}%")
    print(f"  方向准确率:    {val_metrics.get('directional_accuracy', 0)*100:.2f}%")

    print(f"\n💰 验证集回测:")
    print(f"  总收益率:     {val_metrics.get('backtest_total_return', 0)*100:.2f}%")
    print(f"  夏普比率:     {val_metrics.get('backtest_sharpe_ratio', 0):.4f}")

    print(f"\n✓ 模型保存至: ./models/checkpoints/dlinear_simple_features/")
    print(f"✓ 特征数量: 9 (OHLCV + RSI + MACD)")
    print("="*80)

if __name__ == "__main__":
    main()