import os
import json
import argparse
import math


def main(args):
    
    left_sam_images = []
    for dir_name in os.listdir(args.sa1b_path):
        if dir_name[:3] != 'sa_':
            continue
        for file in os.listdir(os.path.join(args.sa1b_path, dir_name)):
            if not file.endswith('jpg'):
                continue
            file_name = file.split('.')[0]
            if not os.path.exists(os.path.join(args.sa1b_path, dir_name, file_name+'.json')):
                continue
            if os.path.exists(os.path.join(args.work_dir, dir_name, file_name+'.json')):
                continue
            left_sam_images.append(os.path.join(args.sa1b_path, dir_name, file))
    
    if not os.path.exists(args.task_file_path):
        os.makedirs(args.task_file_path)

    mini_batch = math.ceil(len(left_sam_images) / args.num_splits)
    for ii in range(args.num_splits):
        with open(os.path.join(args.task_file_path, f"rank{ii}.json"), mode='w') as savef:
            savef.write(json.dumps(left_sam_images[ii*mini_batch:(ii+1)*mini_batch]))


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--num_splits", type=int, default=8)
    parser.add_argument("--sa1b_path", type=str, default="./data/sam/")
    parser.add_argument("--work_dir", type=str, default="./work_dirs/object_level")
    parser.add_argument("--task_file_path", type=str, default="./work_dirs/object_level_task")
    args = parser.parse_args()

    main(args)