import os
import sys
from tqdm import tqdm
from glob import glob

SKIP = [
]

PROCESSED = [
    'ArgoVerse2sync',
    'Kubric4D', 
    'Kubric5D',
    'PD4D',
    'LyftL5',
    'NuScenes',
    'WaymoPerception',
]

OUROBOROS = [
    'DDAD',
    'IODAcore',
    'PDv1',
    'PDv2',
    'PDnvs25',
    'PDnvs89',
]

IRREGULAR = {
    'AgiBotWorldAlpha': 'agibotworld',
    'DROID101': 'droid',
}

DOWNLOADED_ROOT = '/data/cv_downloaded' 
DOWNLOADED_ROOT_PROCESSED = '/data/cv_processed' 
UNIFIED_ROOT = '/data/cv_unified' 
SCRIPT = 'bash scripts/unify_debug.sh'


if __name__ == "__main__":

    datasets = [os.path.basename(d) for d in sorted(glob(f'{DOWNLOADED_ROOT}/*'))]
    datasets_processed = [os.path.basename(d) for d in sorted(glob(f'{DOWNLOADED_ROOT_PROCESSED}/*'))]
    datasets = sorted(datasets + datasets_processed)
    datasets = [d for d  in datasets if d not in SKIP]

    if len(sys.argv) > 1:
        datasets = [d for d in sys.argv[1:] if d in datasets]

    progress = tqdm(datasets, ncols=96)
    for dataset in progress:
        progress.set_description(dataset)
        if os.path.exists(f'{UNIFIED_ROOT}/{dataset}'):
            continue

        if dataset in PROCESSED:
            converter = 'processed'
        elif dataset in OUROBOROS:
            converter = 'ouroboros'
        elif dataset in IRREGULAR:
            converter = IRREGULAR[dataset]
        else: 
            converter = dataset.lower()

        print()
        command = f'{SCRIPT} {converter} {dataset} 1/0 --restart'
        os.system(command)