"""Copies annotations from one file to another."""
import os
from os.path import isdir, join, exists
from glob import glob
import shutil
import numpy as np


if __name__ == "__main__":
    src_dir = "/ssd/pbagad/datasets/CondensedMovies/annotations/"
    
    num_subdirs = 3
    dst_dirs = [f"./annotations_{i}" for i in range(1, num_subdirs + 1)]
    print("Number of destination dirs:", len(dst_dirs))
    
    # load all source files
    src_files = glob(join(src_dir, "*_filtered.json"))
    print("Number of source files:", len(src_files))
    
    # split into num_subdirs
    src_file_splits = np.array_split(np.array(src_files), num_subdirs)
    
    print("Copying")
    # copy to each dst dir
    for src_files, dst_dir in zip(src_file_splits, dst_dirs):
        if not isdir(dst_dir):
            os.makedirs(dst_dir)
        for src_file in src_files:
            dst_file = join(dst_dir, src_file.split("/")[-1])
            # print(f"Copying {src_file} to {dst_file}")
            shutil.copy(src_file, dst_file)