# coding:utf-8
import os
import shutil

def extract_data(source_folder, target_folder,feat_number):
    # 创建目标文件夹
    if not os.path.exists(target_folder):
        os.makedirs(target_folder)

    # 遍历源文件夹中的每个子文件夹
    for subdir in os.listdir(source_folder):
        subdir_path = os.path.join(source_folder, subdir)
        # 获取子文件夹中的前n个文件,os.scandir防止文件过多逐个遍历
        i=0
        if os.path.isdir(subdir_path):
            for entry in os.scandir(subdir_path):
                if i == feat_number:
                    break
                else:
                    i=i+1
                if entry.is_file():
                    file_path = os.path.join(subdir_path, entry.name)
                    print(os.path.join(target_folder, subdir_path.split(os.sep)[-1]))
                    os.makedirs(os.path.join(target_folder, subdir_path.split(os.sep)[-1]), exist_ok=True)
                    shutil.copy(file_path, os.path.join(target_folder, subdir_path.split(os.sep)[-1]))
                    print(f"Extracted {entry.name} from {subdir}")

def main(source_folder, target_folder,feat_number):

    extract_data(source_folder, target_folder,feat_number)


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument("--source_folder", type=str)
    parser.add_argument("--target_folder", type=str)
    parser.add_argument("--feat_number", type=int)

    args = parser.parse_args()

    main(args.source_folder, args.target_folder,args.feat_number)