import os
import yaml
import argparse
from utils import list_s3

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

parser = argparse.ArgumentParser()
parser.add_argument('--dataset', type=str)
parser.add_argument("--filters", type=str, default=None)
parser.add_argument("--subset", nargs='+', type=int, default=None)
parser.add_argument("--subset_oxe", nargs='+', type=int, default=None)
parser.add_argument('--subset_oxe_start', type=str, default='None')
parser.add_argument("--check_only", action='store_true')
args = parser.parse_args()

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

from externals.vidar.vidar.utils.config import read_config
cfg = read_config('custom/prepare/process/config.yaml').datasets.dict[args.dataset]
s3_path = cfg['s3_path'][0]

if args.dataset in ['LBMv12','UMI','RBY']:
    paths = []
    folders1 = list_s3(s3_path)
    for f1 in folders1:
        folders2 = list_s3(f'{s3_path}/{f1}')
        for f2 in folders2:
            folders3 = list_s3(f'{s3_path}/{f1}/{f2}')
            for f3 in folders3:
                paths.append(f'{f1}/{f2}/{f3}')
else:
    paths = list_s3(s3_path)

if args.dataset in ['LBMv12','RBY']:
    paths = [p for p in paths if p.startswith('real/') or p.startswith('sim/')]
elif args.dataset in ['UMI']:
    paths = [p for p in paths if p.startswith('umi/')]

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

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]

if args.subset_oxe is not None:
    args.subset_oxe = f'{args.subset_oxe[0]}-{args.subset_oxe[1]}'
else:
    args.subset_oxe = 'None'

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

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

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

for path in paths:
    script = 'prepare_LBM_check.sh' if args.check_only else 'prepare_LBM.sh'
    command = f'bash custom/utils/process/{script} {args.dataset} {path} {args.subset_oxe} {args.subset_oxe_start}'
    print('##########', command)
    os.system(command)