# 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
        lambda: DepthwiseConv2D(kernel_size,
                                strides=stride,
                                padding='same',
                                activation=activation_fn,
                                use_bias=False), {},
        name="DepthwiseConv2D_%dx%d" % (kernel_size, kernel_size))


# depthwise convolution (followed by relu and batch norm)
def depthwise_conv2d_cell(kernel_size, stride=1):
    return mo.siso_sequential(
        [depthwise_conv2d(kernel_size, stride, 'relu'),
         batch_normalization()])


def separable_conv2d(filters, kernel_size, stride=1, activation_fn='linear'):
    return htfe.siso_tensorflow_eager_module_from_tensorflow_op_fn(
        lambda: SeparableConv2D(filters,
                                kernel_size,
                                strides=stride,
                                padding='same',
                                activation=activation_fn,
                                use_bias=False), {},
        name="SeparableConv2D_%dx%d" % (kernel_size, kernel_size))


# 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'),

# 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

        lambda: DepthwiseConv2D(kernel_size,
                                strides=stride,
                                padding='same',
                                activation=activation_fn,
                                use_bias=False), {},
        name="DepthwiseConv2D_%dx%d" % (kernel_size, kernel_size))


# depthwise convolution (followed by relu and batch norm)
def depthwise_conv2d_cell(kernel_size, stride=1):
    return mo.siso_sequential(
        [depthwise_conv2d(kernel_size, stride, 'relu'),
         batch_normalization()])


def separable_conv2d(filters, kernel_size, stride=1, activation_fn='linear'):
    return htfe.siso_tensorflow_eager_module_from_tensorflow_op_fn(
        lambda: SeparableConv2D(filters,
                                kernel_size,
                                strides=stride,
                                padding='same',
                                activation=activation_fn,
                                use_bias=False), {},
        name="SeparableConv2D_%dx%d" % (kernel_size, kernel_size))


# 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()