#!/usr/bin/env python3

import os
import sys
import shutil
import json
import random

def main():
    dataset_base = '/data/dataset'
    shutil.rmtree(dataset_base + '/pos')
    shutil.rmtree(dataset_base + '/neg')
    os.makedirs(dataset_base + '/pos', exist_ok=True)
    os.makedirs(dataset_base + '/neg', exist_ok=True)
    all_samples = os.listdir('data/samples')
    random.shuffle(all_samples)
    for sample in all_samples[:1000]:
        newname = f'{sample}.jpg'
        md_name = f'data/samples/{sample}/metadata.json'
        pos_or_neg = None
        if 'pos' in sample:
            pos_or_neg = 'pos'
        elif 'neg' in sample:
            pos_or_neg = 'neg'
        elif os.path.exists(md_name):
            with open(md_name, 'r') as f:
                md = json.load(f)
            if 'pos' in md['labels']:
                pos_or_neg = 'pos'
            elif 'neg' in md['labels']:
                pos_or_neg = 'neg'
        if not pos_or_neg:
            continue
        src = f'data/samples/{sample}/full-sbs.jpg'
        if not os.path.exists(src):
            continue
        print(src)
        shutil.copy(src, os.path.join(dataset_base, pos_or_neg, newname))

if __name__ == '__main__':
    main()