import sys
import subprocess 
import argparse
from glob import glob
import yaml

##################################

parser = argparse.ArgumentParser()
parser.add_argument('--dataset', type=str)
parser.add_argument("--filters", type=str, default=None)
parser.add_argument("--buffer", type=str, default='120_p16')
parser.add_argument("--subset", nargs='+', type=int, default=None)
parser.add_argument("--num_proc", type=int, default=16)
parser.add_argument("--delete", action='store_true')
parser.add_argument("--merge", action='store_true')
parser.add_argument('--start_from', type=str, default=None)
parser.add_argument('--only', type=str, default=None)
parser.add_argument("--latents", action='store_true')
args = parser.parse_args()

##################################

if args.latents:
    script = 'python externals/vidar/scripts/datasets/webdataset/create_latent.py'
else:
    script = 'python externals/vidar/scripts/datasets/webdataset/create.py'
cfg_path = 'custom/webdataset/configs'

s3_path = 's3://tri-ml-datasets/cv_datasets/processed'
local_path = '/data/cv_datasets/processed/LBMv12'

recipe = args.dataset
buffer = args.buffer
num_proc = args.num_proc

##################################

paths = []
folders1 = glob(f'{local_path}/*')
for folder1 in folders1:
    folders2 = glob(f'{folder1}/*')
    for folder2 in folders2:
        folders3 = glob(f'{folder2}/*')
        paths.extend(folders3)

##################################

if args.filters is not None:
    with open(args.filters) as stream:
        filters = yaml.safe_load(stream)

    for i, filter in enumerate(filters['filters']):
        filtered_paths = []
        for path in paths:
            valid = False
            for f in filter[f'filter{i}']:
                if f in path:
                    valid = True
                    break
            if valid:
                filtered_paths.append(path)
        paths = filtered_paths

##################################

if args.subset is not None:
    tot, n = args.subset
    paths = paths[n::tot]

##################################

for i in range(len(paths)):
    paths[i] = paths[i].replace(f'{local_path}/', '')

##################################

if args.start_from is not None:
    if args.start_from not in paths:
        print(f"Error: The value of --start_from ('{args.start_from}') is not found in the paths.")
        sys.exit(1)
    idx = paths.index(args.start_from)
    paths = paths[idx:]

if args.only is not None:
    idx = paths.index(args.only)
    paths = [paths[idx]]

##################################

print()
print(f'####################### PROCESSING {len(paths)} datasets ###')
for path in paths:
    print(path)
print(f'####################### PROCESSING {len(paths)} datasets ###')
print()

##################################

for i, path in enumerate(paths):
    command = f'{script} {recipe}'      \
        f' --cfg_path {cfg_path}'       \
        f' --buffer {buffer}'           \
        f' --suffix {path}'             \
        f' --num_proc {num_proc}'       
    if args.delete or (i == 0 and args.start_from is not None):
        command += ' --delete'
    if args.merge:
        command += ' --merge'
    print('##########', command)
    subprocess.call(command, shell=True)


# f' --cameras {camera}'          \
# f' --delete'