import multiprocessing
import os
from tqdm import tqdm

from externals.vidar.scripts.datasets.utils import \
    arg_parser, get_local_files_and_folders, get_splits


args = arg_parser()


def delete_folder(files, n, st, fn):
    progress = range(st, fn)
    if n == 0 or args.multi_bar:
        progress = tqdm(progress, ncols=128)
    for i in progress:
        command = f'rm -rf {files[i]}'
        os.system(command)

def delete_dataset(path):

    files, folders = get_local_files_and_folders(path)
    splits = get_splits(folders, args.num_proc)

    sufs = set([f.split('.')[-1] for f in files])
    for suf in sufs:
        command = f'rm {path}/*.{suf}'
        os.system(command)

    jobs = []
    for i in range(len(splits) - 1):
        process = multiprocessing.Process(
            target=delete_folder,
            args=(folders, i, splits[i], splits[i + 1]),
        )
        jobs.append(process)
    for j in jobs: j.start()
    for j in jobs: j.join()

    command = f'rm -rf {path}'
    os.system(command)

import sys
if __name__ == "__main__":
    paths = sys.argv[1:]
    print('##########################################')
    for path in paths:
        if '/' not in path:
            path = f'/data/cv_datasets/processed/{path}'
        print(f'### !!! DELETING !!! {path}')
        delete_dataset(path)
    print('##########################################')