# Path: fully_supervised/fs_baseline.py
# class FSBaseline(ClassifyFullySupervised, models.MultiModel):
# 
#     def model(self, batch, lr, wd, ema, **kwargs):
#         hwc = [self.dataset.height, self.dataset.width, self.dataset.colors]
#         xt_in = tf.placeholder(tf.float32, [batch] + hwc, 'xt')  # For training
#         x_in = tf.placeholder(tf.float32, [None] + hwc, 'x')
#         l_in = tf.placeholder(tf.int32, [batch], 'labels')
#         wd *= lr
#         classifier = lambda x, **kw: self.classifier(x, **kw, **kwargs).logits
# 
#         x, labels_x = self.augment(xt_in, tf.one_hot(l_in, self.nclass), **kwargs)
#         logits_x = classifier(x, training=True)
# 
#         loss_xe = tf.nn.softmax_cross_entropy_with_logits_v2(labels=labels_x, logits=logits_x)
#         loss_xe = tf.reduce_mean(loss_xe)
#         tf.summary.scalar('losses/xe', loss_xe)
# 
#         ema = tf.train.ExponentialMovingAverage(decay=ema)
#         ema_op = ema.apply(utils.model_vars())
#         ema_getter = functools.partial(utils.getter_ema, ema)
#         post_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS) + [ema_op]
#         post_ops.extend([tf.assign(v, v * (1 - wd)) for v in utils.model_vars('classify') if 'kernel' in v.name])
# 
#         train_op = tf.train.AdamOptimizer(lr).minimize(loss_xe, colocate_gradients_with_ops=True)
#         with tf.control_dependencies([train_op]):
#             train_op = tf.group(*post_ops)
# 
#         return EasyDict(
#             xt=xt_in, x=x_in, label=l_in, train_op=train_op,
#             classify_raw=tf.nn.softmax(classifier(x_in, training=False)),  # No EMA, for debugging.
#             classify_op=tf.nn.softmax(classifier(x_in, getter=ema_getter, training=False)))
# 
# Path: fully_supervised/lib/data.py
# DATASETS = create_datasets
# 
# Path: libml/utils.py
# _GPUS = None
# FLAGS = flags.FLAGS
#         _GPUS = tuple([x.name for x in local_device_protos if x.device_type == 'GPU'])
# class EasyDict(dict):
#     def __init__(self, *args, **kwargs):
# def get_config():
# def setup_main():
# 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 get_gpu():
# def average_gradients(tower_grads):
# def para_list(fn, *args):
# def para_mean(fn, *args):
# def para_cat(fn, *args):
# def combine_dicts(*args):


import os
import tensorflow as tf
from absl import app
from absl import flags
from fully_supervised.fs_baseline import FSBaseline
from fully_supervised.lib.data import DATASETS
from libml import utils
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Mixup fully supervised training.
"""




FLAGS = flags.FLAGS


class FSMixup(FSBaseline):

    def augment(self, x, l, beta, **kwargs):
        del kwargs
        with tf.device('/cpu'):
            mix = tf.distributions.Beta(beta, beta).sample([tf.shape(x)[0], 1, 1, 1])
        mix = tf.maximum(mix, 1 - mix)
        xmix = x * mix + x[::-1] * (1 - mix)
        lmix = l * mix[:, :, 0, 0] + l[::-1] * (1 - mix[:, :, 0, 0])
        return xmix, lmix


def main(argv):

# Path: fully_supervised/fs_baseline.py
# class FSBaseline(ClassifyFullySupervised, models.MultiModel):
# 
#     def model(self, batch, lr, wd, ema, **kwargs):
#         hwc = [self.dataset.height, self.dataset.width, self.dataset.colors]
#         xt_in = tf.placeholder(tf.float32, [batch] + hwc, 'xt')  # For training
#         x_in = tf.placeholder(tf.float32, [None] + hwc, 'x')
#         l_in = tf.placeholder(tf.int32, [batch], 'labels')
#         wd *= lr
#         classifier = lambda x, **kw: self.classifier(x, **kw, **kwargs).logits
# 
#         x, labels_x = self.augment(xt_in, tf.one_hot(l_in, self.nclass), **kwargs)
#         logits_x = classifier(x, training=True)
# 
#         loss_xe = tf.nn.softmax_cross_entropy_with_logits_v2(labels=labels_x, logits=logits_x)
#         loss_xe = tf.reduce_mean(loss_xe)
#         tf.summary.scalar('losses/xe', loss_xe)
# 
#         ema = tf.train.ExponentialMovingAverage(decay=ema)
#         ema_op = ema.apply(utils.model_vars())
#         ema_getter = functools.partial(utils.getter_ema, ema)
#         post_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS) + [ema_op]
#         post_ops.extend([tf.assign(v, v * (1 - wd)) for v in utils.model_vars('classify') if 'kernel' in v.name])
# 
#         train_op = tf.train.AdamOptimizer(lr).minimize(loss_xe, colocate_gradients_with_ops=True)
#         with tf.control_dependencies([train_op]):
#             train_op = tf.group(*post_ops)
# 
#         return EasyDict(
#             xt=xt_in, x=x_in, label=l_in, train_op=train_op,
#             classify_raw=tf.nn.softmax(classifier(x_in, training=False)),  # No EMA, for debugging.
#             classify_op=tf.nn.softmax(classifier(x_in, getter=ema_getter, training=False)))
# 
# Path: fully_supervised/lib/data.py
# DATASETS = create_datasets
# 
# Path: libml/utils.py
# _GPUS = None
# FLAGS = flags.FLAGS
#         _GPUS = tuple([x.name for x in local_device_protos if x.device_type == 'GPU'])
# class EasyDict(dict):
#     def __init__(self, *args, **kwargs):
# def get_config():
# def setup_main():
# 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 get_gpu():
# def average_gradients(tower_grads):
# def para_list(fn, *args):
# def para_mean(fn, *args):
# def para_cat(fn, *args):
# def combine_dicts(*args):


# Path: fully_supervised/fs_mixup.py
import os
import tensorflow as tf
from absl import app
from absl import flags
from fully_supervised.fs_baseline import FSBaseline
from fully_supervised.lib.data import DATASETS
from libml import utils

#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Mixup fully supervised training.
"""




FLAGS = flags.FLAGS


class FSMixup(FSBaseline):

    def augment(self, x, l, beta, **kwargs):
        del kwargs
        with tf.device('/cpu'):
            mix = tf.distributions.Beta(beta, beta).sample([tf.shape(x)[0], 1, 1, 1])
        mix = tf.maximum(mix, 1 - mix)
        xmix = x * mix + x[::-1] * (1 - mix)
        lmix = l * mix[:, :, 0, 0] + l[::-1] * (1 - mix[:, :, 0, 0])
        return xmix, lmix


def main(argv):

    utils.setup_main()