# Path: libml/data.py
# _DATA_CACHE = None
# DATA_DIR = os.environ['ML_DATA']
# FLAGS = flags.FLAGS
# DATASETS = {}
# def record_parse(serialized_example):
# def default_parse(dataset: tf.data.Dataset, parse_fn=record_parse) -> tf.data.Dataset:
# def dataset(filenames: list) -> tf.data.Dataset:
# def memoize(dataset: tf.data.Dataset) -> tf.data.Dataset:
#     def tf_get(index):
#         def get(index):
# def augment_mirror(x):
# def augment_shift(x, w):
# def augment_noise(x, std):
# def compute_mean_std(data: tf.data.Dataset):
#         def iterator():
#     def __init__(self, name, train_labeled, train_unlabeled, test, valid, eval_labeled, eval_unlabeled,
#                  height=32, width=32, colors=3, nclass=10, mean=0, std=1, p_labeled=None, p_unlabeled=None):
#     def creator(cls, name, seed, label, valid, augment, parse_fn=default_parse, do_memoize=True, colors=3,
#                 nclass=10, height=32, width=32, name_suffix=''):
#         def create():
# class DataSet:
# 
# Path: libml/utils.py
# _GPUS = None
# FLAGS = flags.FLAGS
#         _GPUS = tuple([x.name for x in local_device_protos if x.device_type == 'GPU'])
# def get_config():
# def setup_tf():
# def smart_shape(x):
# def ilog2(x):
# def find_latest_checkpoint(dir, glob_term='model.ckpt-*.meta'):
# def get_latest_global_step(dir):
# def get_latest_global_step_in_subdir(dir):
# def getter_ema(ema, getter, name, *args, **kwargs):
# def model_vars(scope=None):
# def gpu(x):
# def get_available_gpus():
# def average_gradients(tower_grads):
# def para_list(fn, *args):
# def para_mean(fn, *args):
# def para_cat(fn, *args):


import hashlib
import itertools
import os
import tensorflow as tf
from absl import app
from absl import flags
from libml import data, utils
from tqdm import trange
DATASETS.update([data.DataSet.creator('svhn', seed, label, valid, lambda x: x, do_memoize=False)
                 for seed, label, valid in
                 itertools.product(range(6), [250, 500, 1000, 2000, 4000], [1, 5000])])
DATASETS.update([data.DataSet.creator('svhn_noextra', seed, label, valid, lambda x: x, do_memoize=False)
                 for seed, label, valid in
                 itertools.product(range(6), [250, 500, 1000, 2000, 4000], [1, 5000])])


def to_byte(d: dict):
    return tf.to_int32(tf.round(127.5 * (d['image'] + 1)))


def collect_hashes(sess, group, data):
    data = data.map(to_byte).batch(FLAGS.batch).prefetch(1).make_one_shot_iterator().get_next()
    hashes = set()
    hasher = hashlib.sha512
    for _ in trange(0, FLAGS.samples, FLAGS.batch, desc='Building hashes for %s' % group, leave=False):
        try:
            batch = sess.run(data)
        except tf.errors.OutOfRangeError:
            break
        for img in batch:
            hashes.add(hasher(img).digest())
    return hashes


def main(argv):
    del argv
    utils.setup_tf()
    dataset = DATASETS[FLAGS.dataset]()

# Path: libml/data.py
# _DATA_CACHE = None
# DATA_DIR = os.environ['ML_DATA']
# FLAGS = flags.FLAGS
# DATASETS = {}
# def record_parse(serialized_example):
# def default_parse(dataset: tf.data.Dataset, parse_fn=record_parse) -> tf.data.Dataset:
# def dataset(filenames: list) -> tf.data.Dataset:
# def memoize(dataset: tf.data.Dataset) -> tf.data.Dataset:
#     def tf_get(index):
#         def get(index):
# def augment_mirror(x):
# def augment_shift(x, w):
# def augment_noise(x, std):
# def compute_mean_std(data: tf.data.Dataset):
#         def iterator():
#     def __init__(self, name, train_labeled, train_unlabeled, test, valid, eval_labeled, eval_unlabeled,
#                  height=32, width=32, colors=3, nclass=10, mean=0, std=1, p_labeled=None, p_unlabeled=None):
#     def creator(cls, name, seed, label, valid, augment, parse_fn=default_parse, do_memoize=True, colors=3,
#                 nclass=10, height=32, width=32, name_suffix=''):
#         def create():
# class DataSet:
# 
# Path: libml/utils.py
# _GPUS = None
# FLAGS = flags.FLAGS
#         _GPUS = tuple([x.name for x in local_device_protos if x.device_type == 'GPU'])
# def get_config():
# def setup_tf():
# def smart_shape(x):
# def ilog2(x):
# def find_latest_checkpoint(dir, glob_term='model.ckpt-*.meta'):
# def get_latest_global_step(dir):
# def get_latest_global_step_in_subdir(dir):
# def getter_ema(ema, getter, name, *args, **kwargs):
# def model_vars(scope=None):
# def gpu(x):
# def get_available_gpus():
# def average_gradients(tower_grads):
# def para_list(fn, *args):
# def para_mean(fn, *args):
# def para_cat(fn, *args):


# Path: scripts/check_split.py
import hashlib
import itertools
import os
import tensorflow as tf
from absl import app
from absl import flags
from libml import data, utils
from tqdm import trange

DATASETS.update([data.DataSet.creator('svhn', seed, label, valid, lambda x: x, do_memoize=False)
                 for seed, label, valid in
                 itertools.product(range(6), [250, 500, 1000, 2000, 4000], [1, 5000])])
DATASETS.update([data.DataSet.creator('svhn_noextra', seed, label, valid, lambda x: x, do_memoize=False)
                 for seed, label, valid in
                 itertools.product(range(6), [250, 500, 1000, 2000, 4000], [1, 5000])])


def to_byte(d: dict):
    return tf.to_int32(tf.round(127.5 * (d['image'] + 1)))


def collect_hashes(sess, group, data):
    data = data.map(to_byte).batch(FLAGS.batch).prefetch(1).make_one_shot_iterator().get_next()
    hashes = set()
    hasher = hashlib.sha512
    for _ in trange(0, FLAGS.samples, FLAGS.batch, desc='Building hashes for %s' % group, leave=False):
        try:
            batch = sess.run(data)
        except tf.errors.OutOfRangeError:
            break
        for img in batch:
            hashes.add(hasher(img).digest())
    return hashes


def main(argv):
    del argv
    utils.setup_tf()
    dataset = DATASETS[FLAGS.dataset]()

    with tf.Session(config=utils.get_config()) as sess: