# Path: deep_architect/contrib/deep_learning_backend/tensorflow_eager_ops.py
# def batch_normalization():
# 
#     def compile_fn(di, dh):
#         bn = tf.layers.BatchNormalization(momentum=.9, epsilon=1e-5)
# 
#         def forward_fn(di, is_training):
#             return {'out': bn(di['in'], training=is_training)}
# 
#         return forward_fn
# 
#     return siso_tensorflow_eager_module('BatchNormalization', compile_fn, {})


from tensorflow.keras.layers import (GlobalAveragePooling2D, Dense, Conv2D,
                                     DepthwiseConv2D, SeparableConv2D,
                                     MaxPooling2D, AveragePooling2D,
                                     Concatenate)
from deep_architect.contrib.deep_learning_backend.tensorflow_eager_ops import batch_normalization
import deep_architect.modules as mo
import deep_architect.hyperparameters as hp
import deep_architect.utils as ut
import deep_architect.helpers.tensorflow_eager_support as htfe


# separable convolution of filters=C channels (followed by relu and batch norm)
def separable_conv2d_cell(filters, kernel_size, stride=1):
    return mo.siso_sequential([
        separable_conv2d(filters, kernel_size, stride, 'relu'),
        batch_normalization()
    ])


# max-pooling
def max_pooling(pool_size, stride=1):
    return htfe.siso_tensorflow_eager_module_from_tensorflow_op_fn(
        lambda: MaxPooling2D(pool_size, strides=stride, padding='same'), {},
        name="MaxPooling2D_%dx%d" % (pool_size, pool_size))


# average-pooling
def average_pooling(pool_size, stride=1):
    return htfe.siso_tensorflow_eager_module_from_tensorflow_op_fn(
        lambda: AveragePooling2D(pool_size, strides=stride, padding='same'), {},
        name="AveragePooling2D_%dx%d" % (pool_size, pool_size))


###############################################################################
# Motif and other functions (level>=1)
###############################################################################


def global_average_pooling():

# Path: deep_architect/contrib/deep_learning_backend/tensorflow_eager_ops.py
# def batch_normalization():
# 
#     def compile_fn(di, dh):
#         bn = tf.layers.BatchNormalization(momentum=.9, epsilon=1e-5)
# 
#         def forward_fn(di, is_training):
#             return {'out': bn(di['in'], training=is_training)}
# 
#         return forward_fn
# 
#     return siso_tensorflow_eager_module('BatchNormalization', compile_fn, {})


# Path: deep_architect/contrib/misc/search_spaces/tensorflow_eager/hierarchical_space.py
from tensorflow.keras.layers import (GlobalAveragePooling2D, Dense, Conv2D,
                                     DepthwiseConv2D, SeparableConv2D,
                                     MaxPooling2D, AveragePooling2D,
                                     Concatenate)
from deep_architect.contrib.deep_learning_backend.tensorflow_eager_ops import batch_normalization
import deep_architect.modules as mo
import deep_architect.hyperparameters as hp
import deep_architect.utils as ut
import deep_architect.helpers.tensorflow_eager_support as htfe



# separable convolution of filters=C channels (followed by relu and batch norm)
def separable_conv2d_cell(filters, kernel_size, stride=1):
    return mo.siso_sequential([
        separable_conv2d(filters, kernel_size, stride, 'relu'),
        batch_normalization()
    ])


# max-pooling
def max_pooling(pool_size, stride=1):
    return htfe.siso_tensorflow_eager_module_from_tensorflow_op_fn(
        lambda: MaxPooling2D(pool_size, strides=stride, padding='same'), {},
        name="MaxPooling2D_%dx%d" % (pool_size, pool_size))


# average-pooling
def average_pooling(pool_size, stride=1):
    return htfe.siso_tensorflow_eager_module_from_tensorflow_op_fn(
        lambda: AveragePooling2D(pool_size, strides=stride, padding='same'), {},
        name="AveragePooling2D_%dx%d" % (pool_size, pool_size))


###############################################################################
# Motif and other functions (level>=1)
###############################################################################


def global_average_pooling():

    return htfe.siso_tensorflow_eager_module_from_tensorflow_op_fn(