import yaml
from utils import list_s3
from tqdm import tqdm

s3_path = 's3://tri-ml-sandbox-16011-us-west-2-datasets/cv_webdatasets'
dataset = 'LBMv12/single_buf[40_p16]-ctx[sequence-single]_cam[scene_left]'
cams = ['scene_left','scene_right','wrist_left_plus','wrist_right_minus']
filters = 'utils/process/filters/filters_sim.yaml'

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

with open(filters) as stream:
    filters = yaml.safe_load(stream)

paths = []
for i, filter in enumerate(filters['filters']):
    if i == 0:
        for f in filter[f'filter{i}']:
            paths.append(f'{s3_path}/{dataset}/{f}')
    else:
        new_paths = []
        for f in filter[f'filter{i}']:
            for p in paths:
                new_paths.append(f'{p}/{f}')
        paths = new_paths        

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

indices = []
for path in tqdm(paths, ncols=128):
    folders1 = list_s3(f'{path}')
    for f1 in folders1:
        indices.append(f'{path}/{f1}/indices/filenames.txt')

cam0 = dataset.split('_cam')[-1][1:-1]
new_indices = []
for ind in indices:
    for cam in cams:
        new_indices.append(ind.replace(cam0, cam))
indices = new_indices

with open('split.txt', 'w') as file:
    for ind in indices:
        file.write(f'[\"{ind.replace(s3_path + "/", "")}\",1],\n')


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

import sys
sys.exit()

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

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

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

print()
print()
print('@@@@@')
for path in paths:
    path = f'[\"{s3_path}/{path}/indices/filenames.txt\",1]'
    print(path)
print()
print()