#!/usr/bin/env python3

import os
import shutil
import json
import argparse

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--num-samples', '-n', type=int, default=100)
    return parser.parse_args()

def main():
    args = parse_args()
    scan_samples = [x for x in os.listdir('data/samples') if 'scan-' in x]
    scan_samples = sorted(scan_samples, reverse=True)
    todo = args.num_samples
    for scan in scan_samples:
        if todo <= 0:
            break
        scan_dir = f'data/samples/{scan}'
        md_file = f'data/samples/{scan}/metadata.json'
        md = json.load(open(md_file))
        scan_id = scan.split('-')[-1]
        targetd = f'/data/emblem/dataset/pos/scans/{scan_id}'
        if os.path.exists(targetd):
            continue
        if 'pos' in md['labels']:
            full_sbs_jpg = f'{scan_dir}/full-sbs.jpg'
            if os.path.exists(full_sbs_jpg):
                os.makedirs(targetd, exist_ok=True)
                print(f'Copying {full_sbs_jpg} => {targetd}/fullqrsbs.jpg')
                shutil.copy(full_sbs_jpg, f'{targetd}/fullqrsbs.jpg')
                print(f'Copying {md_file} => {targetd}/metadata.json')
                shutil.copy(md_file, f'{targetd}/metadata.json')
                todo -= 1

if __name__ == '__main__':
    main()