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

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 models.mlstm_model import MLSTMTransformer as MatrixLSTM
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 train_product_model_with_mlstm(
    product_id,
    model_identifier,
    product_df,
    store_id=None,
    training_mode='product',
    aggregation_method='sum',
    epochs=50,
    sequence_length=LOOK_BACK,
    forecast_horizon=FORECAST_HORIZON,
    model_dir=DEFAULT_MODEL_DIR,
    version=None,
    socketio=None,
    task_id=None,
    continue_training=False,
    progress_callback=None,
    patience=10,
    learning_rate=0.001,
    clip_norm=1.0
):
    """
    使用mLSTM训练产品销售预测模型
    
    参数:
        product_id: 产品ID
        store_id: 店铺ID，为None时使用全局数据
        training_mode: 训练模式 ('product', 'store', 'global')
        aggregation_method: 聚合方法 ('sum', 'mean', 'weighted')
        epochs: 训练轮次
        model_dir: 模型保存目录
        version: 模型版本，如果为None则自动生成
        socketio: Socket.IO实例，用于实时进度推送
        task_id: 任务ID
        continue_training: 是否继续训练
        progress_callback: 进度回调函数，用于多进程训练
    """
    
    # 创建WebSocket进度反馈函数，支持多进程
    def emit_progress(message, progress=None, metrics=None):
        """发送训练进度到前端"""
        progress_data = {
            'task_id': task_id,
            'message': message,
            'timestamp': time.time()
        }
        if progress is not None:
            progress_data['progress'] = progress
        if metrics is not None:
            progress_data['metrics'] = metrics
            
        # 在多进程环境中使用progress_callback
        if progress_callback:
            try:
                progress_callback(progress_data)
            except Exception as e:
                print(f"[mLSTM] 进度回调失败: {e}")
        
        # 在单进程环境中使用socketio
        if socketio and task_id:
            try:
                socketio.emit('training_progress', progress_data, namespace='/training')
            except Exception as e:
                print(f"[mLSTM] WebSocket发送失败: {e}")
                
        print(f"[mLSTM] {message}", flush=True)
        # 强制刷新输出缓冲区
        import sys
        sys.stdout.flush()
        sys.stderr.flush()
    
    emit_progress("开始mLSTM模型训练...")
    
    # 确定版本号
    emit_progress(f"开始训练 mLSTM 模型")
    if version:
        emit_progress(f"使用指定版本: {version}")
    
    # 初始化训练进度管理器（如果还未初始化）
    if socketio and task_id:
        print(f"[mLSTM] 任务 {task_id}: 开始mLSTM训练器", flush=True)
        try:
            # 初始化进度管理器
            if not hasattr(progress_manager, 'training_id') or progress_manager.training_id != task_id:
                progress_manager.start_training(
                    training_id=task_id,
                    product_id=product_id,
                    model_type='mlstm',
                    training_mode=training_mode,
                    total_epochs=epochs,
                    total_batches=0,  # 将在后面设置
                    batch_size=32,    # 默认值
                    total_samples=0   # 将在后面设置
                )
                print(f"[mLSTM] 任务 {task_id}: 进度管理器已初始化", flush=True)
            else:
                print(f"[mLSTM] 任务 {task_id}: 使用现有进度管理器", flush=True)
        except Exception as e:
            print(f"[mLSTM] 任务 {task_id}: 进度管理器初始化失败: {e}", flush=True)

    # 数据现在由调用方传入，不再在此处加载
    if training_mode == 'store' and store_id:
        training_scope = f"店铺 {store_id}"
    elif training_mode == 'global':
        training_scope = f"全局聚合({aggregation_method})"
    else:
        training_scope = "所有店铺"
    
    # 数据量检查
    min_required_samples = sequence_length + forecast_horizon
    if len(product_df) < min_required_samples:
        error_msg = (
            f"❌ 训练数据不足错误\n"
            f"当前配置需要: {min_required_samples} 天数据 (LOOK_BACK={sequence_length} + 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)
        emit_progress(f"训练失败：数据不足 ({len(product_df)}/{min_required_samples} 天)")
        raise ValueError(error_msg)
    
    product_name = product_df['product_name'].iloc[0]
    
    print(f"[mLSTM] 使用mLSTM模型训练产品 '{product_name}' (ID: {product_id}) 的销售预测模型", flush=True)
    print(f"[mLSTM] 训练范围: {training_scope}", flush=True)
    # print(f"[mLSTM] 版本: {version}", flush=True) # Version is now handled by model_manager
    print(f"[mLSTM] 使用设备: {DEVICE}", flush=True)
    print(f"[mLSTM] 模型将保存到目录: {model_dir}", flush=True)
    print(f"[mLSTM] 数据量: {len(product_df)} 条记录", flush=True)
    
    emit_progress(f"训练产品: {product_name} (ID: {product_id}) - {training_scope}")
    
    # 创建特征和目标变量
    features = ['sales', 'weekday', 'month', 'is_holiday', 'is_weekend', 'is_promotion', 'temperature']
    
    print(f"[mLSTM] 开始数据预处理，特征: {features}", flush=True)
    
    # 预处理数据
    X = product_df[features].values
    y = product_df[['sales']].values  # 保持为二维数组
    
    print(f"[mLSTM] 特征矩阵形状: {X.shape}, 目标矩阵形状: {y.shape}", flush=True)
    emit_progress("数据预处理中...")
    
    # 归一化数据
    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)
    
    print(f"[mLSTM] 数据归一化完成", flush=True)
    
    # 划分训练集和测试集（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, sequence_length, forecast_horizon)
    testX, testY = create_dataset(X_test, y_test, sequence_length, forecast_horizon)
    
    # 转换为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)
    
    print(f"[mLSTM] 数据加载器创建完成 - 批次数: {total_batches}, 样本数: {total_samples}", flush=True)
    emit_progress(f"数据加载器准备完成 - 批次数: {total_batches}, 样本数: {total_samples}")
    
    # 初始化mLSTM结合Transformer模型
    input_dim = X_train.shape[1]
    output_dim = forecast_horizon
    hidden_size = 128
    num_heads = 4
    dropout_rate = 0.1
    num_blocks = 3
    embed_dim = 32
    dense_dim = 32
    
    print(f"[mLSTM] 初始化模型 - 输入维度: {input_dim}, 输出维度: {output_dim}", flush=True)
    print(f"[mLSTM] 模型参数 - 隐藏层: {hidden_size}, 注意力头: {num_heads}", flush=True)
    emit_progress(f"初始化mLSTM模型 - 输入维度: {input_dim}, 隐藏层: {hidden_size}")
    
    model = MatrixLSTM(
        num_features=input_dim,
        hidden_size=hidden_size,
        mlstm_layers=2,
        embed_dim=embed_dim,
        dense_dim=dense_dim,
        num_heads=num_heads,
        dropout_rate=dropout_rate,
        num_blocks=num_blocks,
        output_sequence_length=output_dim
    )
    
    print(f"[mLSTM] 模型创建完成", flush=True)
    emit_progress("mLSTM模型初始化完成")
    
    # 如果是继续训练，加载现有模型
    if continue_training and version != 'v1':
        # TODO: Implement continue_training logic with the new model_manager
        pass
    
    # 将模型移动到设备上
    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)

    emit_progress("数据预处理完成，开始模型训练...", progress=10)
    
    # 训练模型
    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
    
    emit_progress(f"开始训练 - 总epoch: {epochs}, 检查点间隔: {checkpoint_interval}, 耐心值: {patience}")
    
    for epoch in range(epochs):
        emit_progress(f"开始训练 Epoch {epoch+1}/{epochs}")
        
        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()
        
        # 计算训练损失
        train_loss = epoch_loss / len(train_loader)
        train_losses.append(train_loss)
        
        # 在测试集上评估
        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()
        
        test_loss = test_loss / len(test_loader)
        test_losses.append(test_loss)
        
        # 更新学习率
        scheduler.step(test_loss)
        
        # 计算总体训练进度
        epoch_progress = ((epoch + 1) / epochs) * 90 + 10  # 10-100% 范围
        
        # 发送训练进度
        current_metrics = {
            'train_loss': train_loss,
            'test_loss': test_loss,
            'epoch': epoch + 1,
            'total_epochs': epochs,
            'learning_rate': optimizer.param_groups[0]['lr']
        }
        
        emit_progress(f"Epoch {epoch+1}/{epochs} 完成 - Train Loss: {train_loss:.4f}, Test Loss: {test_loss:.4f}", 
                     progress=epoch_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_rate': dropout_rate,
                    'num_blocks': num_blocks,
                    'embed_dim': embed_dim,
                    'dense_dim': dense_dim,
                    'mlstm_layers': 2, # 确保这个参数被保存
                    'sequence_length': sequence_length,
                    'forecast_horizon': forecast_horizon,
                    'model_type': 'mlstm'
                },
                'training_info': {
                    'product_id': product_id,
                    'product_name': product_name,
                    'training_mode': training_mode,
                    'store_id': store_id,
                    'aggregation_method': aggregation_method,
                    'training_scope': training_scope,
                    'timestamp': time.time()
                }
            }
            
            # 如果是最佳模型，额外保存一份
            if test_loss < best_loss:
                best_loss = test_loss
                model_manager.save_model(
                    model_data=checkpoint_data,
                    product_id=model_identifier, # 修正：使用唯一的标识符
                    model_type='mlstm',
                    store_id=store_id,
                    training_mode=training_mode,
                    aggregation_method=aggregation_method,
                    product_name=product_name,
                    version='best'
                )
                emit_progress(f"💾 保存最佳模型检查点 (epoch {epoch+1}, test_loss: {test_loss:.4f})")
                epochs_no_improve = 0
            else:
                epochs_no_improve += 1
        
        if (epoch + 1) % 10 == 0:
            print(f"Epoch {epoch+1}/{epochs}, Train Loss: {train_loss:.4f}, Test Loss: {test_loss:.4f}", flush=True)

        # 提前停止逻辑
        if epochs_no_improve >= patience:
            emit_progress(f"连续 {patience} 个epoch测试损失未改善，提前停止训练。")
            break
    
    # 计算训练时间
    training_time = time.time() - start_time
    
    emit_progress("生成损失曲线...", progress=95)
    
    # 确定模型保存目录（支持多店铺）
    if store_id:
        # 为特定店铺创建子目录
        store_model_dir = os.path.join(model_dir, 'mlstm', store_id)
        os.makedirs(store_model_dir, exist_ok=True)
        loss_curve_filename = f"{product_id}_mlstm_{version}_loss_curve.png"
        loss_curve_path = os.path.join(store_model_dir, loss_curve_filename)
    else:
        # 全局模型保存在global目录
        global_model_dir = os.path.join(model_dir, 'mlstm', 'global')
        os.makedirs(global_model_dir, exist_ok=True)
        loss_curve_filename = f"{product_id}_mlstm_{version}_global_loss_curve.png"
        loss_curve_path = os.path.join(global_model_dir, loss_curve_filename)
    
    # 绘制损失曲线并保存到模型目录
    plt.figure(figsize=(10, 6))
    plt.plot(train_losses, label='Training Loss')
    plt.plot(test_losses, label='Test Loss')
    title_suffix = f" - {training_scope}" if store_id else " - 全局模型"
    plt.title(f'mLSTM 模型训练损失曲线 - {product_name} ({version}){title_suffix}')
    plt.xlabel('Epoch')
    plt.ylabel('Loss')
    plt.legend()
    plt.grid(True)
    plt.savefig(loss_curve_path, dpi=300, bbox_inches='tight')
    plt.close()
    
    print(f"损失曲线已保存到: {loss_curve_path}")
    
    emit_progress("模型评估中...", progress=98)
    
    # 评估模型
    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("\n模型评估指标:")
    print(f"MSE: {metrics['mse']:.4f}")
    print(f"RMSE: {metrics['rmse']:.4f}")
    print(f"MAE: {metrics['mae']:.4f}")
    print(f"R²: {metrics['r2']:.4f}")
    print(f"MAPE: {metrics['mape']:.2f}%")
    print(f"训练时间: {training_time:.2f}秒")
    
    emit_progress("保存最终模型...", progress=99)
    
    # 保存最终训练完成的模型（基于最终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_rate': dropout_rate,
            'num_blocks': num_blocks,
            'embed_dim': embed_dim,
            'dense_dim': dense_dim,
            'mlstm_layers': 2, # 确保这个参数被保存
            'sequence_length': sequence_length,
            'forecast_horizon': forecast_horizon,
            'model_type': 'mlstm'
        },
        '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,
            'training_scope': training_scope,
            'timestamp': time.time(),
            'training_completed': True
        }
    }
    
    # 保存最终模型，让 model_manager 自动处理版本号
    final_model_path, final_version = model_manager.save_model(
        model_data=final_model_data,
        product_id=model_identifier, # 修正：使用唯一的标识符
        model_type='mlstm',
        store_id=store_id,
        training_mode=training_mode,
        aggregation_method=aggregation_method,
        product_name=product_name
    )
    
    # 发送训练完成消息
    final_metrics = {
        'mse': metrics['mse'],
        'rmse': metrics['rmse'],
        'mae': metrics['mae'],
        'r2': metrics['r2'],
        'mape': metrics['mape'],
        'training_time': training_time,
        'final_epoch': epochs,
        'model_path': final_model_path,
        'version': final_version
    }
    
    emit_progress(f"✅ mLSTM模型训练完成！版本 {final_version} 已保存", progress=100, metrics=final_metrics)
    
    return model, metrics, epochs, final_model_path