import os
import yaml
import argparse
from utils import list_s3
from tqdm import tqdm

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

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("--s3_path", type=str, default=f's3://tri-ml-sandbox-16011-us-west-2-datasets/cv_datasets/processed')
parser.add_argument("--local_path", type=str, default=f'/data/cv_datasets/processed')
args = parser.parse_args()

args.s3_path = f"{args.s3_path}/{args.dataset}"
args.local_path = f"{args.local_path}/{args.dataset}"

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

def filter_tasks(paths, filters):
    if filters is None:
        return paths
    with open(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
    return paths

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

if args.dataset in ['LBMv12']:
    paths = []
    folders1 = list_s3(args.s3_path)
    folders1 = filter_tasks(folders1, args.filters)
    for f1 in tqdm(folders1):
        folders2 = list_s3(f'{args.s3_path}/{f1}')
        for f2 in folders2:
            folders3 = list_s3(f'{args.s3_path}/{f1}/{f2}')
            for f3 in folders3:
                paths.append(f'{f1}/{f2}/{f3}')
elif args.dataset in ['DROID','BridgeV2']:
    paths = list_s3(args.s3_path)
else:
    raise ValueError('Invalid dataset')

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

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

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

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

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

for path in paths:
    command = f"bash custom/sync/sync_from.sh {args.s3_path}/{path} "
    if args.filters is None:
        command += f"{args.local_path}/{path}"
    else:
        filters = args.filters.split('/')[-1].split('.')[0]
        command += f"{args.local_path}/{filters}/{path}"
    print('##########', command)
    os.system(command)