"""分类任务：使用2020-2025全部数据训练，2026年测试"""

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 ClassificationPreprocessor:
    """分类任务预处理器：预测涨(1)或跌(0)"""

    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):
        """添加技术指标"""
        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=上涨, 0=下跌"""
        df = df.copy()
        df['future_price'] = df['close'].shift(-self.forecast_horizon)
        df['future_return'] = (df['future_price'] - df['close']) / df['close']
        # 二分类：大于0为上涨(1)，否则为下跌(0)
        df['label'] = (df['future_return'] > 0).astype(int)
        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])
            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])

        X, y = [], []
        for i in range(len(scaled_features) - self.lookback):
            X.append(scaled_features[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("分类任务：2020-2025全量训练 → 2026测试")
    print("="*80)

    # ========== 步骤1: 获取训练数据（2020-2025） ==========
    print("\n[1/6] 获取2020-2025年训练数据...")
    fetcher = BinanceDataFetcher()
    df_train = fetcher.fetch_ohlcv(
        symbol='BTC',
        timeframe='1h',
        start_date='2020-01-01',
        end_date='2025-12-31',
        use_cache=True
    )

    print(f"✓ 训练数据量: {len(df_train):,} 行")
    print(f"✓ 时间范围: {df_train['timestamp'].min()} 至 {df_train['timestamp'].max()}")
    print(f"✓ 价格: ${df_train['close'].iloc[0]:,.2f} → ${df_train['close'].iloc[-1]:,.2f}")

    # ========== 步骤2: 预处理训练数据 ==========
    print("\n[2/6] 预处理训练数据（分类任务）...")
    preprocessor = ClassificationPreprocessor(lookback=168, forecast_horizon=1)
    X_train, y_train, processed_train = preprocessor.preprocess(df_train, fit=True)

    print(f"✓ 训练序列: {len(X_train):,}")
    print(f"✓ 特征数量: {len(preprocessor.feature_columns)}")
    print(f"✓ 特征列表: {preprocessor.feature_columns}")

    # 标签分布
    n_up = (y_train == 1).sum()
    n_down = (y_train == 0).sum()
    print(f"\n✓ 标签分布:")
    print(f"   上涨(1): {n_up:,} ({n_up/len(y_train)*100:.1f}%)")
    print(f"   下跌(0): {n_down:,} ({n_down/len(y_train)*100:.1f}%)")

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

    # ========== 步骤3: 创建数据加载器 ==========
    print("\n[3/6] 创建数据加载器...")
    train_dataset = TensorDataset(
        torch.FloatTensor(X_train),
        torch.FloatTensor(y_train)
    )
    # 用10%作为验证集（内部分割）
    train_size = int(0.9 * len(train_dataset))
    val_size = len(train_dataset) - train_size
    train_subset, val_subset = torch.utils.data.random_split(
        train_dataset, [train_size, val_size]
    )

    train_loader = DataLoader(train_subset, batch_size=64, shuffle=True)
    val_loader = DataLoader(val_subset, batch_size=64, shuffle=False)

    print(f"  训练样本: {train_size:,}")
    print(f"  验证样本: {val_size:,}")

    # ========== 步骤4: 创建模型 ==========
    print("\n[4/6] 创建DLinear分类模型...")
    model_config = DLinearConfig(
        input_dim=9,
        lookback=168,
        task='classification',  # 分类任务
        hidden_dim=256,
        dropout=0.1,
        moving_avg_kernel=25
    )
    model = DLinear(model_config)
    print(f"✓ 模型参数: {model.count_parameters():,}")
    print(f"✓ 任务类型: 二分类（预测涨/跌）")
    print(f"✓ 损失函数: BCEWithLogitsLoss")

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

    trainer = Trainer(
        model,
        train_config,
        task='classification',
        pos_weight=pos_weight  # 使用类别权重
    )
    history = trainer.fit(train_loader, val_loader)

    # ========== 步骤6: 保存 ==========
    print("\n[6/6] 保存模型...")
    with open('./models/checkpoints/dlinear_classification_2026/preprocessor.pkl', 'wb') as f:
        pickle.dump(preprocessor, f)

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

    print("\n" + "="*80)
    print("✓ 训练完成！准备在2026年数据上测试...")
    print("="*80)

if __name__ == "__main__":
    main()