import os
import random
import time
import argparse
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.metrics import confusion_matrix, precision_recall_curve, f1_score, precision_score, recall_score, accuracy_score
from tqdm import tqdm

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torch.optim.lr_scheduler import ReduceLROnPlateau

from dataset import AudioDataset
from models import waveform_resnet18, waveform_resnet34, waveform_resnet50, waveform_resnet101, spectrogram_resnet18, spectrogram_resnet34, spectrogram_resnet50, spectrogram_resnet101
import warnings
warnings.filterwarnings("ignore", message="At least one mel filterbank has all zero values")

def parse_args():
    parser = argparse.ArgumentParser(description='训练音频情感分类模型')
    
    # 数据集参数
    parser.add_argument('--data_root', type=str, default='dataset', help='数据集根目录')
    parser.add_argument('--use_mfcc', action='store_true', help='使用MFCC特征而非原始波形')
    parser.add_argument('--val_split', type=float, default=0.2, help='验证集比例')
    parser.add_argument('--num_workers', type=int, default=4, help='数据加载器的工作线程数量')
    
    # 模型参数
    parser.add_argument('--model', type=str, default='resnet18', 
                       choices=['resnet18', 'resnet34', 'resnet50', 'resnet101'],
                       help='模型架构')
    parser.add_argument('--pretrained', action='store_true', help='使用预训练权重')
    
    # 训练参数
    parser.add_argument('--batch_size', type=int, default=32, help='批次大小')
    parser.add_argument('--epochs', type=int, default=30, help='训练轮数')
    parser.add_argument('--lr', type=float, default=0.001, help='学习率')
    parser.add_argument('--weight_decay', type=float, default=1e-5, help='权重衰减')
    parser.add_argument('--optimizer', type=str, default='adam', choices=['adam', 'sgd'], help='优化器选择')
    parser.add_argument('--scheduler', action='store_true', help='是否使用学习率调度器')
    
    # 其他参数
    parser.add_argument('--seed', type=int, default=42, help='随机种子')
    parser.add_argument('--gpu', type=int, default=0, help='使用的GPU ID')
    parser.add_argument('--save_dir', type=str, default='checkpoints', help='模型保存目录')
    
    return parser.parse_args()

def set_seed(seed):
    torch.manual_seed(seed)
    np.random.seed(seed)
    random.seed(seed)
    # 如果有GPU，还需要设置以下代码
    if torch.cuda.is_available():
        torch.cuda.manual_seed(seed)
        torch.cuda.manual_seed_all(seed)
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False

def get_model(args, num_classes):
    if args.model == 'waveform_resnet18':
        return waveform_resnet18(num_classes=num_classes)
    elif args.model == 'waveform_resnet34':
        return waveform_resnet34(num_classes=num_classes)
    elif args.model == 'waveform_resnet50':
        return waveform_resnet50(num_classes=num_classes)
    elif args.model == 'waveform_resnet101':
        return waveform_resnet101(num_classes=num_classes)
    elif args.model == 'spectrogram_resnet18':
        return spectrogram_resnet18(num_classes=num_classes, pretrained=args.pretrained)
    elif args.model == 'spectrogram_resnet34':
        return spectrogram_resnet34(num_classes=num_classes, pretrained=args.pretrained)
    elif args.model == 'spectrogram_resnet50':
        return spectrogram_resnet50(num_classes=num_classes, pretrained=args.pretrained)
    elif args.model == 'spectrogram_resnet101':
        return spectrogram_resnet101(num_classes=num_classes, pretrained=args.pretrained)
    else:
        raise ValueError(f"不支持的模型: {args.model}")

def get_optimizer(args, model):
    if args.optimizer == 'adam':
        return optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay)
    elif args.optimizer == 'sgd':
        return optim.SGD(model.parameters(), lr=args.lr, momentum=0.9, weight_decay=args.weight_decay)
    else:
        raise ValueError(f"不支持的优化器: {args.optimizer}")

def calculate_metrics(outputs, targets, num_classes):
    """计算准确率、精确率、召回率和F1分数"""
    _, preds = torch.max(outputs, 1)
    
    # 转为CPU numpy数组以用于scikit-learn
    preds_np = preds.cpu().numpy()
    targets_np = targets.cpu().numpy()
    
    # 计算总体准确率
    acc = accuracy_score(targets_np, preds_np)
    
    # 计算每个类别的精确率、召回率和F1分数，然后求平均
    precision = precision_score(targets_np, preds_np, average='macro', zero_division=0)
    recall = recall_score(targets_np, preds_np, average='macro', zero_division=0)
    f1 = f1_score(targets_np, preds_np, average='macro', zero_division=0)
    
    return acc, precision, recall, f1

def validate(model, val_loader, criterion, device, num_classes):
    """在验证集上评估模型"""
    model.eval()
    val_loss = 0.0
    all_outputs = []
    all_targets = []
    
    with torch.no_grad():
        for inputs, targets in tqdm(val_loader, desc="验证中", leave=False):
            inputs, targets = inputs.to(device), targets.to(device)
            
            outputs = model(inputs)
            loss = criterion(outputs, targets)
            
            val_loss += loss.item() * inputs.size(0)
            all_outputs.append(outputs)
            all_targets.append(targets)
    
    val_loss /= len(val_loader.dataset)
    
    # 合并所有批次的输出和目标
    all_outputs = torch.cat(all_outputs, dim=0)
    all_targets = torch.cat(all_targets, dim=0)
    
    # 计算评估指标
    acc, precision, recall, f1 = calculate_metrics(all_outputs, all_targets, num_classes)
    
    return val_loss, acc, precision, recall, f1, all_outputs, all_targets

def save_model(model, path):
    """保存模型和配置"""
    os.makedirs(os.path.dirname(path), exist_ok=True)
    torch.save(model.state_dict(), path)
    print(f"模型已保存到 {path}")

def plot_metrics(train_metrics, val_metrics, save_dir):
    """绘制训练和验证指标"""
    metrics = ['loss', 'accuracy', 'precision', 'recall', 'f1']
    
    # 创建保存目录
    os.makedirs(save_dir, exist_ok=True)
    
    for i, metric in enumerate(metrics):
        plt.figure(figsize=(10, 6))
        plt.plot(train_metrics[i], label=f'Train {metric}')
        plt.plot(val_metrics[i], label=f'Validation {metric}')
        plt.xlabel('Epoch')
        plt.ylabel(metric.capitalize())
        plt.title(f'Training and Validation {metric.capitalize()}')
        plt.legend()
        plt.grid(True)
        plt.savefig(os.path.join(save_dir, f'{metric}_plot.png'), dpi=300, bbox_inches='tight')
        plt.close()

def plot_confusion_matrix(true_labels, pred_labels, class_names, save_path):
    """绘制混淆矩阵"""
    cm = confusion_matrix(true_labels, pred_labels)
    plt.figure(figsize=(10, 8))
    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', xticklabels=class_names, yticklabels=class_names)
    plt.xlabel('Predicted')
    plt.ylabel('True')
    plt.title('Confusion Matrix')
    plt.savefig(save_path, dpi=300, bbox_inches='tight')
    plt.close()

def plot_pr_curve(all_outputs, all_targets, num_classes, class_names, save_path):
    """绘制PR曲线"""
    # 将输出转换为概率
    probs = torch.nn.functional.softmax(all_outputs, dim=1).cpu().numpy()
    targets = all_targets.cpu().numpy()
    
    plt.figure(figsize=(12, 10))
    
    # 为每个类别绘制PR曲线
    for i in range(num_classes):
        precision, recall, _ = precision_recall_curve(
            (targets == i).astype(int),
            probs[:, i]
        )
        plt.plot(recall, precision, lw=2, label=f'Class {class_names[i]}')
    
    plt.xlabel('Recall')
    plt.ylabel('Precision')
    plt.title('Precision-Recall Curve')
    plt.legend(loc='best')
    plt.grid(True)
    plt.savefig(save_path, dpi=300, bbox_inches='tight')
    plt.close()

def main():
    args = parse_args()

    # 判断是否使用MFCC特征，如果使用则在模型名前加上spectrogram
    if args.use_mfcc:
        args.model = 'spectrogram_' + args.model
    else:
        args.model = 'waveform_' + args.model

    set_seed(args.seed)
    
    # 设置设备
    device = torch.device(f'cuda:{args.gpu}' if torch.cuda.is_available() else 'cpu')
    print(f"使用设备: {device}")
    
    # 给save_dir添加模型名称、是否使用MFCC特征和日期时间
    args.save_dir = os.path.join(args.save_dir, f"{args.model}_{time.strftime('%Y%m%d_%H%M%S')}")
    print(f"模型保存目录: {args.save_dir}")

    # 创建保存目录
    os.makedirs(args.save_dir, exist_ok=True)
    
    # 加载数据集 - 修改为分别加载训练集和验证集
    train_dir = os.path.join(args.data_root, 'train')
    val_dir = os.path.join(args.data_root, 'val')
    
    train_dataset = AudioDataset(train_dir, use_mfcc=args.use_mfcc)
    val_dataset = AudioDataset(val_dir, use_mfcc=args.use_mfcc)
    
    # 测试集使用验证集
    test_dataset = val_dataset
    
    # 定义类别名称
    class_names = list(train_dataset.label_to_int.keys())
    num_classes = len(class_names)
    print(f"类别数量: {num_classes}")
    print(f"类别名称: {class_names}")
    
    print(f"训练集大小: {len(train_dataset)}")
    print(f"验证集大小: {len(val_dataset)}")
    print(f"测试集大小: {len(test_dataset)}")
    
    # 创建数据加载器
    train_loader = DataLoader(
        train_dataset, batch_size=args.batch_size, shuffle=True, 
        num_workers=args.num_workers, pin_memory=True
    )
    
    val_loader = DataLoader(
        val_dataset, batch_size=args.batch_size, shuffle=False, 
        num_workers=args.num_workers, pin_memory=True
    )
    
    test_loader = DataLoader(
        test_dataset, batch_size=args.batch_size, shuffle=False, 
        num_workers=args.num_workers, pin_memory=True
    )
    
    # 创建模型
    model = get_model(args, num_classes).to(device)
    print(f"模型总参数数量: {sum(p.numel() for p in model.parameters() if p.requires_grad)}")
    
    # 定义损失函数和优化器
    criterion = nn.CrossEntropyLoss()
    optimizer = get_optimizer(args, model)
    
    # 学习率调度器
    scheduler = None
    if args.scheduler:
        scheduler = ReduceLROnPlateau(optimizer, 'min', patience=3, factor=0.5, verbose=True)
    
    # 训练跟踪变量
    best_val_f1 = 0.0
    best_epoch = 0
    best_model_path = ""
    
    # 保存训练和验证指标
    train_losses = []
    train_accs = []
    train_precisions = []
    train_recalls = []
    train_f1s = []
    
    val_losses = []
    val_accs = []
    val_precisions = []
    val_recalls = []
    val_f1s = []
    
    # 记录最佳验证集输出和标签（用于绘图）
    best_val_outputs = None
    best_val_targets = None
    
    # 训练循环
    start_time = time.time()
    for epoch in range(args.epochs):
        model.train()
        train_loss = 0.0
        all_outputs = []
        all_targets = []
        
        train_pbar = tqdm(train_loader, desc=f"Epoch {epoch+1}/{args.epochs}")
        for inputs, targets in train_pbar:
            inputs, targets = inputs.to(device), targets.to(device)
            
            # 前向传播
            optimizer.zero_grad()
            outputs = model(inputs)
            loss = criterion(outputs, targets)
            
            # 反向传播
            loss.backward()
            optimizer.step()
            
            train_loss += loss.item() * inputs.size(0)
            all_outputs.append(outputs)
            all_targets.append(targets)
            
            # 更新进度条
            train_pbar.set_postfix({"Loss": f"{loss.item():.4f}"})
        
        # 计算平均训练损失
        train_loss /= len(train_dataset)
        
        # 计算训练指标
        all_outputs = torch.cat(all_outputs, dim=0)
        all_targets = torch.cat(all_targets, dim=0)
        train_acc, train_precision, train_recall, train_f1 = calculate_metrics(all_outputs, all_targets, num_classes)
        
        # 验证
        val_loss, val_acc, val_precision, val_recall, val_f1, val_outputs, val_targets = validate(
            model, val_loader, criterion, device, num_classes
        )
        
        # 学习率调整
        if scheduler:
            scheduler.step(val_loss)
        
        # 保存指标
        train_losses.append(train_loss)
        train_accs.append(train_acc)
        train_precisions.append(train_precision)
        train_recalls.append(train_recall)
        train_f1s.append(train_f1)
        
        val_losses.append(val_loss)
        val_accs.append(val_acc)
        val_precisions.append(val_precision)
        val_recalls.append(val_recall)
        val_f1s.append(val_f1)
        
        # 打印当前epoch的指标
        print(f"Epoch {epoch+1}/{args.epochs}")
        print(f"Train Loss: {train_loss:.4f} | Val Loss: {val_loss:.4f}")
        print(f"Train Acc: {train_acc:.4f} | Val Acc: {val_acc:.4f}")
        print(f"Train Precision: {train_precision:.4f} | Val Precision: {val_precision:.4f}")
        print(f"Train Recall: {train_recall:.4f} | Val Recall: {val_recall:.4f}")
        print(f"Train F1: {train_f1:.4f} | Val F1: {val_f1:.4f}")
        
        # 保存最佳模型
        if val_f1 > best_val_f1:
            best_val_f1 = val_f1
            best_epoch = epoch
            best_model_path = os.path.join(args.save_dir, f"best.pth")
            save_model(model, best_model_path)
            # 保存最佳验证结果用于后续绘图
            best_val_outputs = val_outputs
            best_val_targets = val_targets
    
    # 训练完成
    total_time = time.time() - start_time
    print(f"训练完成，总耗时: {total_time/60:.2f} 分钟")
    print(f"最佳验证 F1: {best_val_f1:.4f}，在 Epoch {best_epoch+1}")
    
    # 绘制训练和验证指标曲线
    plot_metrics(
        [train_losses, train_accs, train_precisions, train_recalls, train_f1s],
        [val_losses, val_accs, val_precisions, val_recalls, val_f1s],
        os.path.join(args.save_dir, 'plots')
    )
    
    # 加载最佳模型进行测试
    print(f"加载最佳模型进行测试...")
    model.load_state_dict(torch.load(best_model_path))
    test_loss, test_acc, test_precision, test_recall, test_f1, test_outputs, test_targets = validate(
        model, test_loader, criterion, device, num_classes
    )
    
    print(f"测试指标:")
    print(f"Loss: {test_loss:.4f}")
    print(f"Accuracy: {test_acc:.4f}")
    print(f"Precision: {test_precision:.4f}")
    print(f"Recall: {test_recall:.4f}")
    print(f"F1 Score: {test_f1:.4f}")
    
    # 混淆矩阵
    _, test_preds = torch.max(test_outputs, 1)
    plot_confusion_matrix(
        test_targets.cpu().numpy(),
        test_preds.cpu().numpy(),
        class_names,
        os.path.join(args.save_dir, 'plots', 'confusion_matrix.png')
    )
    
    # PR曲线
    plot_pr_curve(
        test_outputs,
        test_targets,
        num_classes,
        class_names,
        os.path.join(args.save_dir, 'plots', 'pr_curve.png')
    )
    
    # 保存测试结果
    results = {
        'test_loss': test_loss,
        'test_accuracy': test_acc,
        'test_precision': test_precision,
        'test_recall': test_recall,
        'test_f1': test_f1,
        'best_epoch': best_epoch + 1,
        'best_val_f1': best_val_f1,
        'train_time': total_time,
        'model': args.model,
        'use_mfcc': args.use_mfcc,
        'pretrained': args.pretrained
    }
    
    # 将结果写入文件
    with open(os.path.join(args.save_dir, 'results.txt'), 'w') as f:
        for key, value in results.items():
            f.write(f"{key}: {value}\n")

if __name__ == "__main__":
    main()