import os
import shutil
import random
import argparse

def split_dataset(source_dir, output_dir, split_ratio=0.8):
    """
    Split the dataset into training and validation sets.

    :param source_dir: Directory containing the dataset with subfolders for each class.
    :param output_dir: Base directory where train and val subdirectories will be created.
    :param split_ratio: Ratio of the training set size to the total dataset size.
    """
    # 创建训练集和验证集目录
    train_dir = os.path.join(output_dir, "train")
    val_dir = os.path.join(output_dir, "val")
    
    if not os.path.exists(train_dir):
        os.makedirs(train_dir)
    if not os.path.exists(val_dir):
        os.makedirs(val_dir)

    for class_dir in os.listdir(source_dir):
        class_path = os.path.join(source_dir, class_dir)
        if os.path.isdir(class_path):
            files = os.listdir(class_path)
            random.shuffle(files)
            split_point = int(len(files) * split_ratio)
            train_files = files[:split_point]
            val_files = files[split_point:]

            train_class_dir = os.path.join(train_dir, class_dir)
            val_class_dir = os.path.join(val_dir, class_dir)
            if not os.path.exists(train_class_dir):
                os.makedirs(train_class_dir)
            if not os.path.exists(val_class_dir):
                os.makedirs(val_class_dir)

            for file in train_files:
                shutil.copy(os.path.join(class_path, file), os.path.join(train_class_dir, file))
            for file in val_files:
                shutil.copy(os.path.join(class_path, file), os.path.join(val_class_dir, file))
    
    # 打印数据集统计信息
    print(f"数据集分割完成:")
    print(f"- 训练集路径: {train_dir}")
    print(f"- 验证集路径: {val_dir}")
    print(f"- 分割比例: {split_ratio:.2f} (训练) / {1-split_ratio:.2f} (验证)")
    
    return train_dir, val_dir

if __name__ == "__main__":
    # 创建命令行参数解析器
    parser = argparse.ArgumentParser(description="将情感语音数据集分割为训练集和验证集")
    parser.add_argument("--source", type=str, default="wav", 
                      help="源数据目录 (默认: wav)")
    parser.add_argument("--output", type=str, default="dataset", 
                      help="输出基础目录，将在其中创建train和val子目录 (默认: dataset)")
    parser.add_argument("--ratio", type=float, default=0.8, 
                      help="训练集比例 (0-1之间的浮点数, 默认: 0.8)")
    parser.add_argument("--seed", type=int, default=42, 
                      help="随机种子 (默认: 无)")
    parser.add_argument("--force", action="store_true", 
                      help="如果目标目录已存在则强制删除")
    
    # 解析命令行参数
    args = parser.parse_args()
    
    # 设置随机种子(如果提供)
    if args.seed is not None:
        random.seed(args.seed)
        print(f"已设置随机种子: {args.seed}")
    
    # 检查分割比例是否有效
    if args.ratio <= 0 or args.ratio >= 1:
        parser.error("分割比例必须在0和1之间")
    
    # 构建实际的训练和验证目录路径
    train_dir = os.path.join(args.output, "train")
    val_dir = os.path.join(args.output, "val")
    
    # 处理目标目录
    if args.force:
        if os.path.exists(train_dir):
            shutil.rmtree(train_dir)
        if os.path.exists(val_dir):
            shutil.rmtree(val_dir)
    else:
        # 检查目标目录是否已存在
        if os.path.exists(train_dir) or os.path.exists(val_dir):
            response = input(f"目标目录 ({train_dir} 或 {val_dir}) 已存在，是否继续并覆盖? (y/n): ")
            if response.lower() != 'y':
                print("操作已取消")
                exit(0)
            else:
                if os.path.exists(train_dir):
                    shutil.rmtree(train_dir)
                if os.path.exists(val_dir):
                    shutil.rmtree(val_dir)
    
    # 执行数据集分割
    split_dataset(args.source, args.output, args.ratio)