# Path: libml/data.py
# class DataSets:
#     def __init__(self, name, train_labeled: DataSet, train_unlabeled: DataSet, test: DataSet, valid: DataSet,
#                  height=32, width=32, colors=3, nclass=10, mean=0, std=1, p_labeled=None, p_unlabeled=None):
#         self.name = name
#         self.train_labeled = train_labeled
#         self.train_unlabeled = train_unlabeled
#         self.test = test
#         self.valid = valid
#         self.height = height
#         self.width = width
#         self.colors = colors
#         self.nclass = nclass
#         self.mean = mean
#         self.std = std
#         self.p_labeled = p_labeled
#         self.p_unlabeled = p_unlabeled
# 
#     @classmethod
#     def creator(cls, name, seed, label, valid, augment, parse_fn=record_parse, do_memoize=False,
#                 nclass=10, colors=3, height=32, width=32):
#         if not isinstance(augment, list):
#             augment = augment(name)
#         fullname = '.%d@%d' % (seed, label)
#         root = os.path.join(DATA_DIR, 'SSL2', name)
# 
#         def create():
#             p_labeled = p_unlabeled = None
# 
#             if FLAGS.p_unlabeled:
#                 sequence = FLAGS.p_unlabeled.split(',')
#                 p_unlabeled = np.array(list(map(float, sequence)), dtype=np.float32)
#                 p_unlabeled /= np.max(p_unlabeled)
# 
#             image_shape = [height, width, colors]
#             train_labeled = DataSet.from_files(
#                 [root + fullname + '-label.tfrecord'], augment[0], parse_fn, image_shape)
#             train_unlabeled = DataSet.from_files(
#                 [root + '-unlabel.tfrecord'], augment[1], parse_fn, image_shape)
#             if do_memoize:
#                 train_labeled = train_labeled.memoize()
#                 train_unlabeled = train_unlabeled.memoize()
# 
#             if FLAGS.whiten:
#                 mean, std = compute_mean_std(train_labeled.concatenate(train_unlabeled))
#             else:
#                 mean, std = 0, 1
# 
#             test_data = DataSet.from_files(
#                 [os.path.join(DATA_DIR, '%s-test.tfrecord' % name)], NOAUGMENT, parse_fn, image_shape=image_shape)
# 
#             return cls(name + '.' + FLAGS.augment + fullname + '-' + str(valid)
#                        + ('/' + FLAGS.p_unlabeled if FLAGS.p_unlabeled else ''),
#                        train_labeled=train_labeled,
#                        train_unlabeled=train_unlabeled.skip(valid),
#                        valid=train_unlabeled.take(valid),
#                        test=test_data,
#                        nclass=nclass, p_labeled=p_labeled, p_unlabeled=p_unlabeled,
#                        height=height, width=width, colors=colors, mean=mean, std=std)
# 
#         return name + fullname + '-' + str(valid), create


import numbers
import numpy as np
import tensorflow as tf
from libml.data import DataSets


def shakeshake(a, b, training):
    if not training:
        return 0.5 * (a + b)
    mu = tf.random_uniform([tf.shape(a)[0]] + [1] * (len(a.shape) - 1), 0, 1)
    mixf = a + mu * (b - a)
    mixb = a + mu[::1] * (b - a)
    return tf.stop_gradient(mixf - mixb) + mixb


class PMovingAverage:
    def __init__(self, name, nclass, buf_size):
        # MEAN aggregation is used by DistributionStrategy to aggregate
        # variable updates across shards
        self.ma = tf.Variable(tf.ones([buf_size, nclass]) / nclass,
                              trainable=False,
                              name=name,
                              aggregation=tf.VariableAggregation.MEAN)

    def __call__(self):
        v = tf.reduce_mean(self.ma, axis=0)
        return v / tf.reduce_sum(v)

    def update(self, entry):
        entry = tf.reduce_mean(entry, axis=0)
        return tf.assign(self.ma, tf.concat([self.ma[1:], [entry]], axis=0))


class PData:

# Path: libml/data.py
# class DataSets:
#     def __init__(self, name, train_labeled: DataSet, train_unlabeled: DataSet, test: DataSet, valid: DataSet,
#                  height=32, width=32, colors=3, nclass=10, mean=0, std=1, p_labeled=None, p_unlabeled=None):
#         self.name = name
#         self.train_labeled = train_labeled
#         self.train_unlabeled = train_unlabeled
#         self.test = test
#         self.valid = valid
#         self.height = height
#         self.width = width
#         self.colors = colors
#         self.nclass = nclass
#         self.mean = mean
#         self.std = std
#         self.p_labeled = p_labeled
#         self.p_unlabeled = p_unlabeled
# 
#     @classmethod
#     def creator(cls, name, seed, label, valid, augment, parse_fn=record_parse, do_memoize=False,
#                 nclass=10, colors=3, height=32, width=32):
#         if not isinstance(augment, list):
#             augment = augment(name)
#         fullname = '.%d@%d' % (seed, label)
#         root = os.path.join(DATA_DIR, 'SSL2', name)
# 
#         def create():
#             p_labeled = p_unlabeled = None
# 
#             if FLAGS.p_unlabeled:
#                 sequence = FLAGS.p_unlabeled.split(',')
#                 p_unlabeled = np.array(list(map(float, sequence)), dtype=np.float32)
#                 p_unlabeled /= np.max(p_unlabeled)
# 
#             image_shape = [height, width, colors]
#             train_labeled = DataSet.from_files(
#                 [root + fullname + '-label.tfrecord'], augment[0], parse_fn, image_shape)
#             train_unlabeled = DataSet.from_files(
#                 [root + '-unlabel.tfrecord'], augment[1], parse_fn, image_shape)
#             if do_memoize:
#                 train_labeled = train_labeled.memoize()
#                 train_unlabeled = train_unlabeled.memoize()
# 
#             if FLAGS.whiten:
#                 mean, std = compute_mean_std(train_labeled.concatenate(train_unlabeled))
#             else:
#                 mean, std = 0, 1
# 
#             test_data = DataSet.from_files(
#                 [os.path.join(DATA_DIR, '%s-test.tfrecord' % name)], NOAUGMENT, parse_fn, image_shape=image_shape)
# 
#             return cls(name + '.' + FLAGS.augment + fullname + '-' + str(valid)
#                        + ('/' + FLAGS.p_unlabeled if FLAGS.p_unlabeled else ''),
#                        train_labeled=train_labeled,
#                        train_unlabeled=train_unlabeled.skip(valid),
#                        valid=train_unlabeled.take(valid),
#                        test=test_data,
#                        nclass=nclass, p_labeled=p_labeled, p_unlabeled=p_unlabeled,
#                        height=height, width=width, colors=colors, mean=mean, std=std)
# 
#         return name + fullname + '-' + str(valid), create


# Path: libml/layers.py
import numbers
import numpy as np
import tensorflow as tf
from libml.data import DataSets



def shakeshake(a, b, training):
    if not training:
        return 0.5 * (a + b)
    mu = tf.random_uniform([tf.shape(a)[0]] + [1] * (len(a.shape) - 1), 0, 1)
    mixf = a + mu * (b - a)
    mixb = a + mu[::1] * (b - a)
    return tf.stop_gradient(mixf - mixb) + mixb


class PMovingAverage:
    def __init__(self, name, nclass, buf_size):
        # MEAN aggregation is used by DistributionStrategy to aggregate
        # variable updates across shards
        self.ma = tf.Variable(tf.ones([buf_size, nclass]) / nclass,
                              trainable=False,
                              name=name,
                              aggregation=tf.VariableAggregation.MEAN)

    def __call__(self):
        v = tf.reduce_mean(self.ma, axis=0)
        return v / tf.reduce_sum(v)

    def update(self, entry):
        entry = tf.reduce_mean(entry, axis=0)
        return tf.assign(self.ma, tf.concat([self.ma[1:], [entry]], axis=0))


class PData:

    def __init__(self, dataset: DataSets):