# Path: selfsup/util.py
# class DummyDict(object):
#     def __init__(self):
#         pass
#     def __getitem__(self, item):
#         return DummyDict()
#     def __setitem__(self, item, value):
#         return DummyDict()
#     def get(self, item, default=None):
#         if default is None:
#             return DummyDict()
#         else:
#             return default
# 
# Path: selfsup/ops.py
# def max_pool(x, size, stride=None, name=None, info=DummyDict(), padding='SAME'):
# def avg_pool(x, size, stride=None, name=None, info=DummyDict(), padding='SAME'):
# def dropout(x, drop_prob, phase_test=None, name=None, info=DummyDict()):
# def scale(x, name=None, value=1.0):
# def inner(x, channels, info=DummyDict(), stddev=None,
#           activation=tf.nn.relu, name=None):
# def atrous_avg_pool(value, size, rate, padding, name=None, info=DummyDict()):
# def conv(x, channels, size=3, strides=1, activation=tf.nn.relu, name=None, padding='SAME',
#          info=DummyDict(), output_shape=None):
# def upconv(x, channels, size=3, strides=1, output_shape=None, activation=tf.nn.relu, name=None, padding='SAME',
#          info=DummyDict()):
#             W = tf.get_variable('weights', shape, dtype=tf.float32,
#                                 initializer=W_init)
#         W = tf.get_variable(name + '/weights', shape, dtype=tf.float32,
#                             initializer=W_init)
#         W = tf.get_variable(name + '/weights', shape, dtype=tf.float32,
#                             initializer=W_init)
# 
# Path: selfsup/model/alex_bn.py
# def batch_norm(z, global_step, phase_test, name, bn_mean=None, bn_var=None):
#     mm, vv = extra.moments(z, list(range(z.get_shape().ndims-1)), keep_dims=False, name=name + '_moments')
# 
#     beta = 0.0
#     gamma = 1.0
# 
#     sh = mm.get_shape().as_list()[-1:]
#     if bn_mean is None and bn_var is None:
#         bn_mean, bn_var = init_batch_norm_vars(name, sh)
# 
#     alpha0 = 0.999
#     N = 1000
#     alpha = tf.to_float(tf.minimum(global_step, N) / N  * alpha0)
#     def mean_var_train():
#         apply_op_mm = tf.assign(bn_mean, bn_mean * alpha + mm * (1 - alpha))
#         apply_op_vv = tf.assign(bn_var, bn_var * alpha + vv * (1 - alpha))
# 
#         with tf.control_dependencies([apply_op_mm, apply_op_vv]):
#             return tf.identity(mm), tf.identity(vv)
#             #return tf.identity(bn_mean), tf.identity(bn_var)
# 
#     def mean_var_test():
#         return bn_mean, bn_var
# 
#     mean, var = tf.cond(tf.logical_not(phase_test),
#                         mean_var_train,
#                         mean_var_test)
# 
#     z = tf.nn.batch_normalization(z, mean, var, beta, gamma, 1e-5)
#     return z


import tensorflow as tf
from selfsup.util import DummyDict
from selfsup import ops
from .alex_bn import batch_norm
        def bn(z, name):
            return batch_norm(z, global_step=global_step, phase_test=phase_test, name=name)
    else:
        def bn(z, name):
            return z

    if len(y.get_shape().as_list()) == 2:
        y = tf.expand_dims(tf.expand_dims(y, 1), 1)

    def check(name):
        return name in LAYERS[LAYERS.index(to_name):LAYERS.index(from_name)+1]

    if check('fc7'):
        sh = [BATCH_SIZE, 1, 1, 4096]
        y = ops.upconv(y, sh[-1], size=1, strides=1, info=info, activation=None, name='upfc7', output_shape=sh)
        y = bn(y, 'upfc7')
        info['activations']['pre:upfc7'] = y
        y = tf.nn.relu(y)
        info['activations']['upfc7'] = y

    if check('fc6'):
        sh = [BATCH_SIZE, 1, 1, 4096]
        y = ops.upconv(y, sh[-1], size=1, strides=1, info=info, activation=None, name='upfc6', output_shape=sh)
        y = bn(y, 'upfc6')
        info['activations']['pre:upfc6'] = y
        y = tf.nn.relu(y)
        info['activations']['upfc6'] = y

    if check('conv5'):
        sh = [BATCH_SIZE, 6, 6, 256]

# Path: selfsup/util.py
# class DummyDict(object):
#     def __init__(self):
#         pass
#     def __getitem__(self, item):
#         return DummyDict()
#     def __setitem__(self, item, value):
#         return DummyDict()
#     def get(self, item, default=None):
#         if default is None:
#             return DummyDict()
#         else:
#             return default
# 
# Path: selfsup/ops.py
# def max_pool(x, size, stride=None, name=None, info=DummyDict(), padding='SAME'):
# def avg_pool(x, size, stride=None, name=None, info=DummyDict(), padding='SAME'):
# def dropout(x, drop_prob, phase_test=None, name=None, info=DummyDict()):
# def scale(x, name=None, value=1.0):
# def inner(x, channels, info=DummyDict(), stddev=None,
#           activation=tf.nn.relu, name=None):
# def atrous_avg_pool(value, size, rate, padding, name=None, info=DummyDict()):
# def conv(x, channels, size=3, strides=1, activation=tf.nn.relu, name=None, padding='SAME',
#          info=DummyDict(), output_shape=None):
# def upconv(x, channels, size=3, strides=1, output_shape=None, activation=tf.nn.relu, name=None, padding='SAME',
#          info=DummyDict()):
#             W = tf.get_variable('weights', shape, dtype=tf.float32,
#                                 initializer=W_init)
#         W = tf.get_variable(name + '/weights', shape, dtype=tf.float32,
#                             initializer=W_init)
#         W = tf.get_variable(name + '/weights', shape, dtype=tf.float32,
#                             initializer=W_init)
# 
# Path: selfsup/model/alex_bn.py
# def batch_norm(z, global_step, phase_test, name, bn_mean=None, bn_var=None):
#     mm, vv = extra.moments(z, list(range(z.get_shape().ndims-1)), keep_dims=False, name=name + '_moments')
# 
#     beta = 0.0
#     gamma = 1.0
# 
#     sh = mm.get_shape().as_list()[-1:]
#     if bn_mean is None and bn_var is None:
#         bn_mean, bn_var = init_batch_norm_vars(name, sh)
# 
#     alpha0 = 0.999
#     N = 1000
#     alpha = tf.to_float(tf.minimum(global_step, N) / N  * alpha0)
#     def mean_var_train():
#         apply_op_mm = tf.assign(bn_mean, bn_mean * alpha + mm * (1 - alpha))
#         apply_op_vv = tf.assign(bn_var, bn_var * alpha + vv * (1 - alpha))
# 
#         with tf.control_dependencies([apply_op_mm, apply_op_vv]):
#             return tf.identity(mm), tf.identity(vv)
#             #return tf.identity(bn_mean), tf.identity(bn_var)
# 
#     def mean_var_test():
#         return bn_mean, bn_var
# 
#     mean, var = tf.cond(tf.logical_not(phase_test),
#                         mean_var_train,
#                         mean_var_test)
# 
#     z = tf.nn.batch_normalization(z, mean, var, beta, gamma, 1e-5)
#     return z


# Path: selfsup/model/alex_transposed.py
import tensorflow as tf
from selfsup.util import DummyDict
from selfsup import ops
from .alex_bn import batch_norm

        def bn(z, name):
            return batch_norm(z, global_step=global_step, phase_test=phase_test, name=name)
    else:
        def bn(z, name):
            return z

    if len(y.get_shape().as_list()) == 2:
        y = tf.expand_dims(tf.expand_dims(y, 1), 1)

    def check(name):
        return name in LAYERS[LAYERS.index(to_name):LAYERS.index(from_name)+1]

    if check('fc7'):
        sh = [BATCH_SIZE, 1, 1, 4096]
        y = ops.upconv(y, sh[-1], size=1, strides=1, info=info, activation=None, name='upfc7', output_shape=sh)
        y = bn(y, 'upfc7')
        info['activations']['pre:upfc7'] = y
        y = tf.nn.relu(y)
        info['activations']['upfc7'] = y

    if check('fc6'):
        sh = [BATCH_SIZE, 1, 1, 4096]
        y = ops.upconv(y, sh[-1], size=1, strides=1, info=info, activation=None, name='upfc6', output_shape=sh)
        y = bn(y, 'upfc6')
        info['activations']['pre:upfc6'] = y
        y = tf.nn.relu(y)
        info['activations']['upfc6'] = y

    if check('conv5'):
        sh = [BATCH_SIZE, 6, 6, 256]

        y = ops.upconv(y, sh[-1], size=6, strides=2, info=info, activation=None, name='upconv5_pre', output_shape=sh, padding='VALID')