# 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:


import itertools
import tensorflow as tf
from absl import flags
from libml.data import DataSet, augment_cifar10, augment_svhn, augment_stl10
flags.DEFINE_integer('nu', 2, 'Number of augmentations for class-consistency.')
FLAGS = flags.FLAGS


def stack_augment(augment):
    def func(x):
        xl = [augment(x) for _ in range(FLAGS.nu)]

        return dict(image=tf.stack([x['image'] for x in xl]),
                    label=tf.stack([x['label'] for x in xl]))

    return func


DATASETS = {}
DATASETS.update([DataSet.creator('cifar10', seed, label, valid, [augment_cifar10, stack_augment(augment_cifar10)])
                 for seed, label, valid in
                 itertools.product(range(6), [250, 500, 1000, 2000, 4000, 8000], [1, 5000])])
DATASETS.update(
    [DataSet.creator('cifar100', seed, label, valid, [augment_cifar10, stack_augment(augment_cifar10)], nclass=100)
     for seed, label, valid in
     itertools.product(range(6), [10000], [1, 5000])])
DATASETS.update([DataSet.creator('stl10', seed, label, valid, [augment_stl10, stack_augment(augment_stl10)], height=96,
                                 width=96, do_memoize=False)
                 for seed, label, valid in
                 itertools.product(range(6), [1000, 5000], [1, 500])])
DATASETS.update([DataSet.creator('svhn', seed, label, valid, [augment_svhn, stack_augment(augment_svhn)],
                                 do_memoize=False)
                 for seed, label, valid in
                 itertools.product(range(6), [250, 500, 1000, 2000, 4000, 8000], [1, 5000])])

# 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/data_pair.py
import itertools
import tensorflow as tf
from absl import flags
from libml.data import DataSet, augment_cifar10, augment_svhn, augment_stl10

flags.DEFINE_integer('nu', 2, 'Number of augmentations for class-consistency.')
FLAGS = flags.FLAGS


def stack_augment(augment):
    def func(x):
        xl = [augment(x) for _ in range(FLAGS.nu)]

        return dict(image=tf.stack([x['image'] for x in xl]),
                    label=tf.stack([x['label'] for x in xl]))

    return func


DATASETS = {}
DATASETS.update([DataSet.creator('cifar10', seed, label, valid, [augment_cifar10, stack_augment(augment_cifar10)])
                 for seed, label, valid in
                 itertools.product(range(6), [250, 500, 1000, 2000, 4000, 8000], [1, 5000])])
DATASETS.update(
    [DataSet.creator('cifar100', seed, label, valid, [augment_cifar10, stack_augment(augment_cifar10)], nclass=100)
     for seed, label, valid in
     itertools.product(range(6), [10000], [1, 5000])])
DATASETS.update([DataSet.creator('stl10', seed, label, valid, [augment_stl10, stack_augment(augment_stl10)], height=96,
                                 width=96, do_memoize=False)
                 for seed, label, valid in
                 itertools.product(range(6), [1000, 5000], [1, 500])])
DATASETS.update([DataSet.creator('svhn', seed, label, valid, [augment_svhn, stack_augment(augment_svhn)],
                                 do_memoize=False)
                 for seed, label, valid in
                 itertools.product(range(6), [250, 500, 1000, 2000, 4000, 8000], [1, 5000])])

DATASETS.update([DataSet.creator('svhn_noextra', seed, label, valid, [augment_svhn, stack_augment(augment_svhn)],