# """
# 将大规模数据集均匀划分为指定数量的子数据集，每个子数据集的样本数差异不超过1。
# 支持从嵌套目录加载 Hugging Face 格式数据集，合并后均匀切分并保存。
# """

# import argparse
# import os
# from os import path

# from tqdm import tqdm
# from datasets import Dataset as HFDataset, concatenate_datasets

# parser = argparse.ArgumentParser(
#     description="Split a large dataset into a specified number of evenly-sized subsets."
# )
# parser.add_argument(
#     "--input_path",
#     type=str,
#     required=True,
#     help="Path to the input dataset (in Hugging Face format, supports nested directories).",
# )
# parser.add_argument(
#     "--output_dir",
#     type=str,
#     required=True,
#     help="Directory to save the split datasets.",
# )
# parser.add_argument(
#     "--num_subsets",
#     type=int,
#     required=True,
#     help="Number of subsets to split the dataset into (must be a positive integer).",
# )
# parser.add_argument(
#     "--shuffle",
#     action="store_true",
#     default=True,
#     help="Whether to shuffle the dataset before splitting (default: True).",
# )
# parser.add_argument(
#     "--no_shuffle",
#     dest="shuffle",
#     action="store_false",
#     help="Do not shuffle the dataset before splitting.",
# )
# parser.add_argument(
#     "--seed",
#     type=int,
#     default=42,
#     help="Random seed for shuffling the dataset (default: 42).",
# )

# args = parser.parse_args()

# if __name__ == "__main__":
#     # 创建输出目录
#     os.makedirs(args.output_dir, exist_ok=True)

#     # ---------------------- 1. 收集所有数据集文件夹路径 ----------------------
#     print("Collecting dataset folders...")
#     dataset_folders = []
#     for data_type in os.listdir(args.input_path):
#         data_type_path = path.join(args.input_path, data_type)
#         if path.isdir(data_type_path):
#             for folder in os.listdir(data_type_path):
#                 folder_path = path.join(data_type_path, folder)
#                 if path.isdir(folder_path):
#                     dataset_folders.append(folder_path)

#     if not dataset_folders:
#         raise ValueError(f"No valid dataset folders found in {args.input_path}")

#     # ---------------------- 2. 加载并合并所有数据集 ----------------------
#     print(f"Loading {len(dataset_folders)} datasets...")
#     data_list = []
#     for folder_path in tqdm(dataset_folders, desc="Loading datasets"):
#         try:
#             dataset = HFDataset.load_from_disk(folder_path)
#             data_list.append(dataset)
#         except Exception as e:
#             print(f"Warning: Failed to load {folder_path}, error: {e}")

#     if not data_list:
#         raise ValueError("No datasets were successfully loaded.")

#     print("Concatenating datasets...")
#     datasets = concatenate_datasets(data_list)
#     total_data = len(datasets)
#     print(f"Total dataset size: {total_data}")

#     # ---------------------- 3. 参数校验（核心修改点1：直接用 num_subsets） ----------------------
#     if args.num_subsets <= 0:
#         raise ValueError("--num_subsets must be a positive integer")
#     if args.num_subsets > total_data:
#         print(
#             f"Warning: num_subsets ({args.num_subsets}) > total_data ({total_data}), adjusting num_subsets to {total_data}"
#         )
#         args.num_subsets = total_data

#     # ---------------------- 4. 可选：打乱数据集 ----------------------
#     if args.shuffle:
#         print(f"Shuffling dataset with seed {args.seed}...")
#         datasets = datasets.shuffle(seed=args.seed)

#     # ---------------------- 5. 均匀划分逻辑（核心修改点2：基于 num_subsets 计算） ----------------------
#     num_subsets = args.num_subsets
#     k = total_data // num_subsets  # 每个子集的基础样本数
#     r = total_data % num_subsets  # 需要多1个样本的子集数量

#     print(f"\nEvenly splitting into {num_subsets} subsets:")
#     print(f"  - {r} subsets with {k+1} samples each")
#     print(f"  - {num_subsets - r} subsets with {k} samples each")

#     # ---------------------- 6. 拆分并保存每个子集 ----------------------
#     current_idx = 0
#     for subset_idx in range(num_subsets):
#         subset_size = k + 1 if subset_idx < r else k
#         start_idx = current_idx
#         end_idx = current_idx + subset_size

#         # 用 select 方法确保返回 Dataset 对象
#         subdataset = datasets.select(range(start_idx, end_idx))

#         # 保存子集（文件名对齐，方便后续按顺序读取）
#         output_path = path.join(args.output_dir, f"dataset_subset_{subset_idx:04d}")
#         subdataset.save_to_disk(output_path)

#         print(f"Saved subset {subset_idx:04d} to {output_path} (size: {subset_size})")

#         current_idx = end_idx

#     print("\nAll subsets saved successfully!")


"""
多进程版：将大规模数据集均匀划分为指定数量的子数据集。
优化点：
1. 多进程并行加载原始数据集
2. 多进程并行保存划分后的子集
3. 保留均匀划分和打乱功能
"""

import argparse
import os
from os import path
import multiprocessing as mp
from functools import partial

from tqdm import tqdm
from datasets import Dataset as HFDataset, concatenate_datasets


# ---------------------- 多进程辅助函数 ----------------------
def load_single_dataset(folder_path):
    """单个进程加载一个数据集文件夹"""
    try:
        dataset = HFDataset.load_from_disk(folder_path)
        return dataset
    except Exception as e:
        print(f"Warning: Failed to load {folder_path}, error: {e}")
        return None


def save_single_subset(subset_info, full_dataset):
    """
    单个进程保存一个子集
    subset_info: (subset_idx, start_idx, end_idx, output_path)
    """
    subset_idx, start_idx, end_idx, output_path = subset_info
    try:
        # 选择当前子集的数据
        subdataset = full_dataset.select(range(start_idx, end_idx))
        # 保存到磁盘
        subdataset.save_to_disk(output_path)
        return (subset_idx, end_idx - start_idx, None)  # (索引, 样本数, 错误信息)
    except Exception as e:
        return (subset_idx, 0, str(e))


# ---------------------- 主逻辑 ----------------------
if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Multi-process split large dataset into evenly-sized subsets."
    )
    parser.add_argument(
        "--input_path",
        type=str,
        required=True,
        help="Path to the input dataset (in Hugging Face format, supports nested directories).",
    )
    parser.add_argument(
        "--output_dir",
        type=str,
        required=True,
        help="Directory to save the split datasets.",
    )
    parser.add_argument(
        "--num_subsets",
        type=int,
        required=True,
        help="Number of subsets to split the dataset into (must be a positive integer).",
    )
    parser.add_argument(
        "--shuffle",
        action="store_true",
        default=True,
        help="Whether to shuffle the dataset before splitting (default: True).",
    )
    parser.add_argument(
        "--no_shuffle",
        dest="shuffle",
        action="store_false",
        help="Do not shuffle the dataset before splitting.",
    )
    parser.add_argument(
        "--seed",
        type=int,
        default=42,
        help="Random seed for shuffling the dataset (default: 42).",
    )
    parser.add_argument(
        "--num_workers",
        type=int,
        default=mp.cpu_count(),
        help=f"Number of processes to use (default: CPU count = {mp.cpu_count()}).",
    )

    args = parser.parse_args()
    os.makedirs(args.output_dir, exist_ok=True)

    # ---------------------- 1. 收集所有数据集文件夹路径 ----------------------
    print("Collecting dataset folders...")
    dataset_folders = []
    for data_type in os.listdir(args.input_path):
        data_type_path = path.join(args.input_path, data_type)
        if path.isdir(data_type_path):
            for folder in os.listdir(data_type_path):
                folder_path = path.join(data_type_path, folder)
                if path.isdir(folder_path):
                    dataset_folders.append(folder_path)

    if not dataset_folders:
        raise ValueError(f"No valid dataset folders found in {args.input_path}")

    # ---------------------- 2. 多进程并行加载原始数据集 ----------------------
    print(
        f"Loading {len(dataset_folders)} datasets with {args.num_workers} processes..."
    )
    with mp.Pool(processes=args.num_workers) as pool:
        data_list = list(
            tqdm(
                pool.imap(load_single_dataset, dataset_folders),
                total=len(dataset_folders),
                desc="Loading datasets",
            )
        )

    # 过滤掉加载失败的数据集
    data_list = [d for d in data_list if d is not None]
    if not data_list:
        raise ValueError("No datasets were successfully loaded.")

    # ---------------------- 3. 合并数据集并打乱 ----------------------
    print("Concatenating datasets...")
    datasets = concatenate_datasets(data_list)
    total_data = len(datasets)
    print(f"Total dataset size: {total_data}")

    # 参数校验
    if args.num_subsets <= 0:
        raise ValueError("--num_subsets must be a positive integer")
    if args.num_subsets > total_data:
        print(
            f"Warning: num_subsets ({args.num_subsets}) > total_data ({total_data}), adjusting num_subsets to {total_data}"
        )
        args.num_subsets = total_data

    # 打乱数据集
    if args.shuffle:
        print(f"Shuffling dataset with seed {args.seed}...")
        datasets = datasets.shuffle(seed=args.seed)

    # ---------------------- 4. 计算子集划分信息 ----------------------
    num_subsets = args.num_subsets
    k = total_data // num_subsets  # 每个子集的基础样本数
    r = total_data % num_subsets  # 需要多1个样本的子集数量

    print(f"\nEvenly splitting into {num_subsets} subsets:")
    print(f"  - {r} subsets with {k+1} samples each")
    print(f"  - {num_subsets - r} subsets with {k} samples each")

    # 生成每个子集的信息：(索引, 起始位置, 结束位置, 保存路径)
    subset_info_list = []
    current_idx = 0
    for subset_idx in range(num_subsets):
        subset_size = k + 1 if subset_idx < r else k
        start_idx = current_idx
        end_idx = current_idx + subset_size
        output_path = path.join(args.output_dir, f"dataset_subset_{subset_idx:04d}")
        subset_info_list.append((subset_idx, start_idx, end_idx, output_path))
        current_idx = end_idx

    for item in subset_info_list:
        print(item)

    # import sys
    # sys.exit()

    # ---------------------- 5. 多进程并行保存子集 ----------------------
    print(f"\nSaving subsets with {args.num_workers} processes...")
    # 使用partial将full_dataset绑定到save_single_subset函数
    save_func = partial(save_single_subset, full_dataset=datasets)

    with mp.Pool(processes=args.num_workers) as pool:
        results = list(
            tqdm(
                pool.imap(save_func, subset_info_list),
                total=len(subset_info_list),
                desc="Saving subsets",
            )
        )

    # ---------------------- 6. 处理保存结果 ----------------------
    success_count = 0
    for subset_idx, subset_size, error in results:
        if error is None:
            print(f"Saved subset {subset_idx:04d} (size: {subset_size})")
            success_count += 1
        else:
            print(f"Failed to save subset {subset_idx:04d}: {error}")

    if success_count == num_subsets:
        print("\nAll subsets saved successfully!")
    else:
        print(f"\nSaved {success_count}/{num_subsets} subsets successfully.")
