#!/usr/bin/env python3


import math
import os
import random

import braceexpand
import webdataset as wds

DEFAULT_CATS_FILE = os.path.join(os.path.dirname(__file__), '..', 'configs', 'places2-categories_157.txt')

def is_good_key(key, cats):
    return any(c in key for c in cats)


def main(args):
    if args.categories == 'nofilter':
        good_categories = None
    else:
        with open(args.categories, 'r') as f:
            good_categories = set(line.strip().split(' ')[0] for line in f if line.strip())

    all_input_files = list(braceexpand.braceexpand(args.infile))
    chunk_size = int(math.ceil(len(all_input_files) / args.n_read_streams))

    input_iterators = [iter(wds.Dataset(all_input_files[start : start + chunk_size]).shuffle(args.shuffle_buffer))
                       for start in range(0, len(all_input_files), chunk_size)]
    output_datasets = [wds.ShardWriter(args.outpattern.format(i)) for i in range(args.n_write_streams)]

    good_readers = list(range(len(input_iterators)))
    step_i = 0
    good_samples = 0
    bad_samples = 0
    while len(good_readers) > 0:
        if step_i % args.print_freq == 0:
            print(f'Iterations done {step_i}; readers alive {good_readers}; good samples {good_samples}; bad samples {bad_samples}')

        step_i += 1

        ri = random.choice(good_readers)
        try:
            sample = next(input_iterators[ri])
        except StopIteration:
            good_readers = list(set(good_readers) - {ri})
            continue

        if good_categories is not None and not is_good_key(sample['__key__'], good_categories):
            bad_samples += 1
            continue

        wi = random.randint(0, args.n_write_streams - 1)
        output_datasets[wi].write(sample)
        good_samples += 1


if __name__ == '__main__':
    import argparse

    aparser = argparse.ArgumentParser()
    aparser.add_argument('--categories', type=str, default=DEFAULT_CATS_FILE)
    aparser.add_argument('--shuffle-buffer', type=int, default=10000)
    aparser.add_argument('--n-read-streams', type=int, default=10)
    aparser.add_argument('--n-write-streams', type=int, default=10)
    aparser.add_argument('--print-freq', type=int, default=1000)
    aparser.add_argument('infile', type=str)
    aparser.add_argument('outpattern', type=str)

    main(aparser.parse_args())