import os
import shutil
from pathlib import Path

# ================== 配置 ==================
SOURCE_DIR = "/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/1/class/class0/3"  # 原始图片文件夹
OUTPUT_A = "/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/1/class/class0/2"  # a 子文件夹
OUTPUT_B = "/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/1/class/class0/4"  # b 子文件夹

# 支持的图片扩展名
IMG_EXTENSIONS = {'.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.webp'}

# ================== 创建输出目录 ==================
os.makedirs(OUTPUT_A, exist_ok=True)
os.makedirs(OUTPUT_B, exist_ok=True)

# ================== 获取并排序图片 ==================
source_path = Path(SOURCE_DIR)
if not source_path.exists():
    raise FileNotFoundError(f"源文件夹不存在: {SOURCE_DIR}")

# 获取所有图片文件，并按名称排序（确保顺序可重现）
image_files = sorted([
    f for f in source_path.iterdir()
    if f.is_file() and f.suffix.lower() in IMG_EXTENSIONS
])

if not image_files:
    print("⚠️ 源文件夹中没有找到图片！")
    exit()

print(f"共找到 {len(image_files)} 张图片，开始分配...")

# ================== 交替分配到 a 和 b ==================
for idx, img_path in enumerate(image_files):
    if idx % 2 == 0:
        # 偶数索引（0, 2, 4...）→ a
        dest = os.path.join(OUTPUT_A, img_path.name)
    else:
        # 奇数索引（1, 3, 5...）→ b
        dest = os.path.join(OUTPUT_B, img_path.name)

    shutil.move(img_path, dest)  # 移动原来数据
    print(f"[{idx + 1}/{len(image_files)}] {img_path.name} → {'a' if idx % 2 == 0 else 'b'}")

print("\n✅ 分配完成！")
print(f"a 文件夹: {OUTPUT_A}")
print(f"b 文件夹: {OUTPUT_B}")