ShopTRAINING/server/trainers/transformer_trainer.py

"""
药店销售预测系统 - Transformer模型训练函数
"""

import os
import time
import pandas as pd
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from sklearn.preprocessing import MinMaxScaler
import matplotlib.pyplot as plt
from tqdm import tqdm
from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score

from models.transformer_model import TimeSeriesTransformer
from utils.data_utils import create_dataset, PharmacyDataset
from utils.multi_store_data_utils import get_store_product_sales_data, aggregate_multi_store_data
from utils.visualization import plot_loss_curve
from analysis.metrics import evaluate_model
from core.config import (
    DEVICE, DEFAULT_MODEL_DIR, LOOK_BACK, FORECAST_HORIZON
)
from utils.training_progress import progress_manager
from utils.model_manager import model_manager

def save_checkpoint(checkpoint_data: dict, epoch_or_label, path_info: dict):
    """
    保存训练检查点
    
    Args:
        checkpoint_data: 检查点数据
        epoch_or_label: epoch编号或标签（如'best', 'final', 50）
        path_info (dict): 包含所有路径信息的字典
    """
    if epoch_or_label == 'best':
        # 使用由 ModelPathManager 直接提供的最佳检查点路径
        checkpoint_path = path_info['best_checkpoint_path']
    else:
        # 使用 epoch 检查点模板生成路径
        template = path_info.get('epoch_checkpoint_template')
        if not template:
            raise ValueError("路径信息 'path_info' 中缺少 'epoch_checkpoint_template'。")
        checkpoint_path = template.format(N=epoch_or_label)

    # 保存检查点
    torch.save(checkpoint_data, checkpoint_path)
    print(f"[Transformer] 检查点已保存: {checkpoint_path}", flush=True)
    
    return checkpoint_path

def train_product_model_with_transformer(
    product_id,
    product_df=None,
    store_id=None,
    training_mode='product',
    aggregation_method='sum',
    epochs=50,
    model_dir=DEFAULT_MODEL_DIR, # 将被 path_info 替代
    version=None, # 将被 path_info 替代
    socketio=None,
    task_id=None,
    continue_training=False,
    path_info=None, # 新增参数
    patience=10,
    learning_rate=0.001,
    clip_norm=1.0
):
    """
    使用Transformer模型训练产品销售预测模型
    
    参数:
    product_id: 产品ID
    epochs: 训练轮次
    model_dir: 模型保存目录，默认使用配置中的DEFAULT_MODEL_DIR
    version: 指定版本号，如果为None则自动生成
    socketio: WebSocket对象，用于实时反馈
    task_id: 训练任务ID
    continue_training: 是否继续训练现有模型
    
    返回:
    model: 训练好的模型
    metrics: 模型评估指标
    version: 实际使用的版本号
    """
    
    if not path_info:
        raise ValueError("train_product_model_with_transformer 需要 'path_info' 参数。")
        
    version = path_info['version']

    # WebSocket进度反馈函数
    def emit_progress(message, progress=None, metrics=None):
        """发送训练进度到前端"""
        if socketio and task_id:
            data = {
                'task_id': task_id,
                'message': message,
                'timestamp': time.time()
            }
            if progress is not None:
                data['progress'] = progress
            if metrics is not None:
                data['metrics'] = metrics
            socketio.emit('training_progress', data, namespace='/training')
        print(f"[{time.strftime('%H:%M:%S')}] {message}", flush=True)
        # 强制刷新输出缓冲区
        import sys
        sys.stdout.flush()
        sys.stderr.flush()
    
    emit_progress(f"开始Transformer模型训练... 版本 v{version}")
    
    # 获取训练进度管理器实例
    try:
        from utils.training_progress import progress_manager
    except ImportError:
        # 如果无法导入，创建一个空的管理器以避免错误
        class DummyProgressManager:
            def set_stage(self, *args, **kwargs): pass
            def start_training(self, *args, **kwargs): pass
            def start_epoch(self, *args, **kwargs): pass
            def update_batch(self, *args, **kwargs): pass
            def finish_epoch(self, *args, **kwargs): pass
            def finish_training(self, *args, **kwargs): pass
        progress_manager = DummyProgressManager()
    
    # 如果没有传入product_df，则根据训练模式加载数据
    if product_df is None:
        from utils.multi_store_data_utils import load_multi_store_data, get_store_product_sales_data, aggregate_multi_store_data
        
        try:
            if training_mode == 'store' and store_id:
                # 加载特定店铺的数据
                product_df = get_store_product_sales_data(
                    store_id,
                    product_id,
                    'pharmacy_sales_multi_store.csv'
                )
                training_scope = f"店铺 {store_id}"
            elif training_mode == 'global':
                # 聚合所有店铺的数据
                product_df = aggregate_multi_store_data(
                    product_id,
                    aggregation_method=aggregation_method,
                    file_path='pharmacy_sales_multi_store.csv'
                )
                training_scope = f"全局聚合({aggregation_method})"
            else:
                # 默认：加载所有店铺的产品数据
                product_df = load_multi_store_data('pharmacy_sales_multi_store.csv', product_id=product_id)
                training_scope = "所有店铺"
        except ValueError as e:
            if "No objects to concatenate" in str(e):
                err_msg = f"聚合数据失败 (product: {product_id}, store: {store_id}, mode: {training_mode}): 没有找到可聚合的数据。"
                emit_progress(err_msg)
                # 在这种情况下，我们不能继续，所以抛出异常
                raise ValueError(err_msg) from e
            # 对于其他 ValueError，也打印并重新抛出
            emit_progress(f"数据加载时发生值错误: {e}")
            raise e
        except Exception as e:
            emit_progress(f"多店铺数据加载失败: {e}, 尝试后备方案...")
            # 后备方案：尝试原始数据
            try:
                df = pd.read_excel('pharmacy_sales.xlsx')
                product_df = df[df['product_id'] == product_id].sort_values('date')
                training_scope = "原始数据"
                emit_progress("成功从 'pharmacy_sales.xlsx' 加载后备数据。")
            except Exception as fallback_e:
                emit_progress(f"后备数据加载失败: {fallback_e}")
                raise fallback_e from e
    else:
        # 如果传入了product_df，直接使用
        if training_mode == 'store' and store_id:
            training_scope = f"店铺 {store_id}"
        elif training_mode == 'global':
            training_scope = f"全局聚合({aggregation_method})"
        else:
            training_scope = "所有店铺"
    
    if product_df.empty:
        raise ValueError(f"产品 {product_id} 没有可用的销售数据")
    
    # 数据量检查
    min_required_samples = LOOK_BACK + FORECAST_HORIZON
    if len(product_df) < min_required_samples:
        error_msg = (
            f"❌ 训练数据不足错误\n"
            f"当前配置需要: {min_required_samples} 天数据 (LOOK_BACK={LOOK_BACK} + FORECAST_HORIZON={FORECAST_HORIZON})\n"
            f"实际数据量: {len(product_df)} 天\n"
            f"产品ID: {product_id}, 训练模式: {training_mode}\n"
            f"建议解决方案:\n"
            f"1. 生成更多数据: uv run generate_multi_store_data.py\n"
            f"2. 调整配置参数: 减小 LOOK_BACK 或 FORECAST_HORIZON\n"
            f"3. 使用全局训练模式聚合更多数据"
        )
        print(error_msg)
        raise ValueError(error_msg)
    
    product_df = product_df.sort_values('date')
    product_name = product_df['product_name'].iloc[0]
    
    print(f"[Transformer] 训练产品 '{product_name}' (ID: {product_id}) 的销售预测模型", flush=True)
    print(f"[Device] 使用设备: {DEVICE}", flush=True)
    print(f"[Model] 模型将保存到: {path_info['base_dir']}", flush=True)
    
    # 创建特征和目标变量
    features = ['sales', 'weekday', 'month', 'is_holiday', 'is_weekend', 'is_promotion', 'temperature']
    
    # 设置数据预处理阶段
    progress_manager.set_stage("data_preprocessing", 0)
    emit_progress("数据预处理中...")
    
    # 预处理数据
    X = product_df[features].values
    y = product_df[['sales']].values  # 保持为二维数组
    
    # 归一化数据
    scaler_X = MinMaxScaler(feature_range=(0, 1))
    scaler_y = MinMaxScaler(feature_range=(0, 1))
    
    X_scaled = scaler_X.fit_transform(X)
    y_scaled = scaler_y.fit_transform(y)
    
    progress_manager.set_stage("data_preprocessing", 40)
    
    # 划分训练集和测试集（80% 训练，20% 测试）
    train_size = int(len(X_scaled) * 0.8)
    X_train, X_test = X_scaled[:train_size], X_scaled[train_size:]
    y_train, y_test = y_scaled[:train_size], y_scaled[train_size:]
    
    # 创建时间序列数据
    trainX, trainY = create_dataset(X_train, y_train, LOOK_BACK, FORECAST_HORIZON)
    testX, testY = create_dataset(X_test, y_test, LOOK_BACK, FORECAST_HORIZON)
    
    progress_manager.set_stage("data_preprocessing", 70)
    
    # 转换为PyTorch的Tensor
    trainX_tensor = torch.Tensor(trainX)
    trainY_tensor = torch.Tensor(trainY)
    testX_tensor = torch.Tensor(testX)
    testY_tensor = torch.Tensor(testY)
    
    # 创建数据加载器
    train_dataset = PharmacyDataset(trainX_tensor, trainY_tensor)
    test_dataset = PharmacyDataset(testX_tensor, testY_tensor)
    
    batch_size = 32
    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
    test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)
    
    # 更新进度管理器的批次信息
    total_batches = len(train_loader)
    total_samples = len(train_dataset)
    progress_manager.total_batches_per_epoch = total_batches
    progress_manager.batch_size = batch_size
    progress_manager.total_samples = total_samples
    
    progress_manager.set_stage("data_preprocessing", 100)
    emit_progress("数据预处理完成，开始模型训练...")
    
    # 初始化Transformer模型
    input_dim = X_train.shape[1]
    output_dim = FORECAST_HORIZON
    hidden_size = 64
    num_heads = 4
    dropout_rate = 0.1
    num_layers = 3
    
    model = TimeSeriesTransformer(
        num_features=input_dim,
        d_model=hidden_size,
        nhead=num_heads,
        num_encoder_layers=num_layers,
        dim_feedforward=hidden_size * 2,
        dropout=dropout_rate,
        output_sequence_length=output_dim,
        seq_length=LOOK_BACK,
        batch_size=batch_size
    )
    
    # 将模型移动到设备上
    model = model.to(DEVICE)
    
    criterion = nn.MSELoss()
    optimizer = optim.Adam(model.parameters(), lr=learning_rate)
    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'min', patience=patience // 2, factor=0.5)
    
    # 训练模型
    train_losses = []
    test_losses = []
    start_time = time.time()
    
    # 配置检查点保存
    checkpoint_interval = max(1, epochs // 10)  # 每10%进度保存一次，最少每1个epoch
    best_loss = float('inf')
    epochs_no_improve = 0
    
    progress_manager.set_stage("model_training", 0)
    emit_progress(f"开始训练 - 总epoch: {epochs}, 检查点间隔: {checkpoint_interval}, 耐心值: {patience}")
    
    for epoch in range(epochs):
        # 开始新的轮次
        progress_manager.start_epoch(epoch)
        
        model.train()
        epoch_loss = 0
        
        for batch_idx, (X_batch, y_batch) in enumerate(train_loader):
            X_batch, y_batch = X_batch.to(DEVICE), y_batch.to(DEVICE)
            
            # 确保目标张量有正确的形状
            # 前向传播
            outputs = model(X_batch)
            loss = criterion(outputs, y_batch)
            
            # 反向传播和优化
            optimizer.zero_grad()
            loss.backward()
            if clip_norm:
                torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=clip_norm)
            optimizer.step()
            
            epoch_loss += loss.item()
            
            # 更新批次进度
            if batch_idx % 5 == 0 or batch_idx == len(train_loader) - 1:
                current_lr = optimizer.param_groups[0]['lr']
                progress_manager.update_batch(batch_idx, loss.item(), current_lr)
        
        # 计算训练损失
        train_loss = epoch_loss / len(train_loader)
        train_losses.append(train_loss)
        
        # 设置验证阶段
        progress_manager.set_stage("validation", 0)
        
        # 在测试集上评估
        model.eval()
        test_loss = 0
        with torch.no_grad():
            for batch_idx, (X_batch, y_batch) in enumerate(test_loader):
                X_batch, y_batch = X_batch.to(DEVICE), y_batch.to(DEVICE)
                
                # 确保目标张量有正确的形状
                outputs = model(X_batch)
                loss = criterion(outputs, y_batch)
                test_loss += loss.item()
                
                # 更新验证进度
                if batch_idx % 3 == 0 or batch_idx == len(test_loader) - 1:
                    val_progress = (batch_idx / len(test_loader)) * 100
                    progress_manager.set_stage("validation", val_progress)
        
        test_loss = test_loss / len(test_loader)
        test_losses.append(test_loss)
        
        # 更新学习率
        scheduler.step(test_loss)
        
        # 完成当前轮次
        progress_manager.finish_epoch(train_loss, test_loss)
        
        # 发送训练进度
        if (epoch + 1) % 5 == 0 or epoch == epochs - 1:
            progress = ((epoch + 1) / epochs) * 100
            current_metrics = {
                'train_loss': train_loss,
                'test_loss': test_loss,
                'epoch': epoch + 1,
                'total_epochs': epochs
            }
            emit_progress(f"Epoch {epoch+1}/{epochs}, Train Loss: {train_loss:.4f}, Test Loss: {test_loss:.4f}", 
                         progress=progress, metrics=current_metrics)
        
        # 定期保存检查点
        if (epoch + 1) % checkpoint_interval == 0 or epoch == epochs - 1:
            checkpoint_data = {
                'epoch': epoch + 1,
                'model_state_dict': model.state_dict(),
                'optimizer_state_dict': optimizer.state_dict(),
                'train_loss': train_loss,
                'test_loss': test_loss,
                'train_losses': train_losses,
                'test_losses': test_losses,
                'scaler_X': scaler_X,
                'scaler_y': scaler_y,
                'config': {
                    'input_dim': input_dim,
                    'output_dim': output_dim,
                    'hidden_size': hidden_size,
                    'num_heads': num_heads,
                    'dropout': dropout_rate,
                    'num_layers': num_layers,
                    'sequence_length': LOOK_BACK,
                    'forecast_horizon': FORECAST_HORIZON,
                    'model_type': 'transformer'
                },
                'training_info': {
                    'product_id': product_id,
                    'product_name': product_name,
                    'training_mode': training_mode,
                    'store_id': store_id,
                    'aggregation_method': aggregation_method,
                    'timestamp': time.time()
                }
            }
            
            # 检查是否为最佳模型
            is_best = test_loss < best_loss
            if is_best:
                best_loss = test_loss
                epochs_no_improve = 0
                # 保存最佳模型检查点
                save_checkpoint(checkpoint_data, 'best', path_info)
                emit_progress(f"💾 保存最佳模型检查点 (epoch {epoch+1}, test_loss: {test_loss:.4f})")
            else:
                epochs_no_improve += 1

            # 保存定期的epoch检查点（如果不是最佳模型，或者即时是最佳也保存一份epoch版本）
            save_checkpoint(checkpoint_data, epoch + 1, path_info)
            emit_progress(f"💾 保存训练检查点 epoch_{epoch+1}")
        
        if (epoch + 1) % 10 == 0:
            print(f"📊 Epoch {epoch+1}/{epochs}, 训练损失: {train_loss:.4f}, 测试损失: {test_loss:.4f}", flush=True)

        # 提前停止逻辑
        if epochs_no_improve >= patience:
            emit_progress(f"连续 {patience} 个epoch测试损失未改善，提前停止训练。")
            break
    
    # 计算训练时间
    training_time = time.time() - start_time
    
    # 设置模型保存阶段
    progress_manager.set_stage("model_saving", 0)
    emit_progress("训练完成，正在保存模型...")
    
    # 绘制损失曲线并保存到模型目录
    loss_curve_path = path_info['loss_curve_path']
    plot_loss_curve(
        train_losses,
        test_losses,
        product_name,
        'Transformer',
        save_path=loss_curve_path
    )
    print(f"📈 损失曲线已保存到: {loss_curve_path}", flush=True)
    
    # 评估模型
    model.eval()
    with torch.no_grad():
        test_pred = model(testX_tensor.to(DEVICE)).cpu().numpy()
        test_true = testY
        
    # 反归一化预测结果和真实值
    test_pred_inv = scaler_y.inverse_transform(test_pred)
    test_true_inv = scaler_y.inverse_transform(test_true)
    
    # 计算评估指标
    metrics = evaluate_model(test_true_inv, test_pred_inv)
    metrics['training_time'] = training_time
    
    # 打印评估指标
    print(f"\n📊 模型评估指标:", flush=True)
    print(f"   MSE: {metrics['mse']:.4f}", flush=True)
    print(f"   RMSE: {metrics['rmse']:.4f}", flush=True)
    print(f"   MAE: {metrics['mae']:.4f}", flush=True)
    print(f"   R²: {metrics['r2']:.4f}", flush=True)
    print(f"   MAPE: {metrics['mape']:.2f}%", flush=True)
    print(f"   ⏱️  训练时间: {training_time:.2f}秒", flush=True)
    
    # 保存最终训练完成的模型（基于最终epoch）
    final_model_data = {
        'epoch': epochs,  # 最终epoch
        'model_state_dict': model.state_dict(),
        'optimizer_state_dict': optimizer.state_dict(),
        'train_loss': train_losses[-1],
        'test_loss': test_losses[-1],
        'train_losses': train_losses,
        'test_losses': test_losses,
        'scaler_X': scaler_X,
        'scaler_y': scaler_y,
        'config': {
            'input_dim': input_dim,
            'output_dim': output_dim,
            'hidden_size': hidden_size,
            'num_heads': num_heads,
            'dropout': dropout_rate,
            'num_layers': num_layers,
            'sequence_length': LOOK_BACK,
            'forecast_horizon': FORECAST_HORIZON,
            'model_type': 'transformer'
        },
        'metrics': metrics,
        'loss_curve_path': loss_curve_path,
        'training_info': {
            'product_id': product_id,
            'product_name': product_name,
            'training_mode': training_mode,
            'store_id': store_id,
            'aggregation_method': aggregation_method,
            'timestamp': time.time(),
            'training_completed': True
        }
    }
    
    progress_manager.set_stage("model_saving", 50)
    
    # 检查模型性能是否达标
    # 移除R2检查，始终保存模型
    if metrics:
        # 保存最终模型
        final_model_path = path_info['model_path']
        torch.save(final_model_data, final_model_path)
        progress_manager.set_stage("model_saving", 100)
        emit_progress(f"模型已保存到 {final_model_path}")
        print(f"💾 模型已保存到 {final_model_path}", flush=True)
    else:
        final_model_path = None
        print(f"[Transformer] 训练过程中未生成评估指标，不保存最终模型。", flush=True)
    
    # 准备最终返回的指标
    final_metrics = {
        'mse': metrics['mse'],
        'rmse': metrics['rmse'],
        'mae': metrics['mae'],
        'r2': metrics['r2'],
        'mape': metrics['mape'],
        'training_time': training_time,
        'final_epoch': epochs
    }

    if final_model_path:
        emit_progress(f"✅ Transformer模型训练完成！", progress=100, metrics=final_metrics)
    else:
        emit_progress(f"❌ Transformer模型训练失败：性能不达标", progress=100, metrics={'error': '模型性能不佳'})
    
    return model, metrics, epochs, final_model_path