import os
import shutil
import math

def split_folder_by_extension(input_folder, output_base_folder, parts=5, move=False, extensions=('.tiff', '.jpg')):
    """
    将指定文件夹中的文件（支持指定扩展名）按比例分成若干份保存到子文件夹中。

    参数:
    - input_folder: 原始文件夹路径
    - output_base_folder: 输出的分组文件夹根路径
    - parts: 要分成几份
    - move: 是否移动文件（True 为移动，False 为复制）
    - extensions: 要处理的文件扩展名组成的元组，例如 ('.tiff', '.jpg')
    """

    # 获取符合扩展名的所有文件
    files = [f for f in os.listdir(input_folder) if f.lower().endswith(extensions)]
    total_files = len(files)

    if total_files == 0:
        print(f"No files found with extensions {extensions}")
        return

    # 排序确保一致性
    files.sort()

    # 每组的文件数
    files_per_group = math.ceil(total_files / parts)

    print(f"Total files: {total_files}")
    print(f"Splitting into {parts} parts, {files_per_group} files per part.")
    print(f"Extensions included: {extensions}")

    for i in range(parts):
        part_folder = os.path.join(output_base_folder, f"part_{i}")
        os.makedirs(part_folder, exist_ok=True)

        start_idx = i * files_per_group
        end_idx = min(start_idx + files_per_group, total_files)

        for filename in files[start_idx:end_idx]:
            src_path = os.path.join(input_folder, filename)
            dst_path = os.path.join(part_folder, filename)

            if move:
                shutil.move(src_path, dst_path)
            else:
                shutil.copy2(src_path, dst_path)

        print(f"Part {i}: {end_idx - start_idx} files saved to {part_folder}")


# ✅ 示例调用
if __name__ == '__main__':
    # input_folder = r'G:\python_ws_g\data\pcd2color_result\color_jpg'
    # output_base_folder = r'G:\python_ws_g\data\average'
    # input_folder = r'\\192.168.50.222\share\zyh\5月彩色钢板数据汇总\515下午 57\pcd2color_result\color_jpg'
    input_folder = r'/data/share/zyh/master_dataset/circle/huayan_circle/day02/color'
    # output_base_folder = r'\\192.168.50.222\share\zyh\5月彩色钢板数据汇总\515下午 57\average'
    output_base_folder = r'/data/share/zyh/master_dataset/circle/huayan_circle/day02/average'

    # 调用函数，自定义扩展名和分几份
    split_folder_by_extension(
        input_folder=input_folder,
        output_base_folder=output_base_folder,
        parts=8,
        move=False,
        extensions=('.tiff', '.jpg', '.png')  # 支持 .tiff 和 .jpg 文件
    )