"""
EfficientDet object detector implementation
"""
import tensorflow as tf

from itertools import count

from .detector_utils import boxes_from_output, non_maximum_suppression


def BackboneStem(filters, name=None):
    """EfficientNet stem block
    """
    def backbone_stem(x):
        x = tf.keras.layers.Conv2D(filters, 3, 2, padding='same', use_bias=False, name=f'{name}.conv2d')(x)
        x = tf.keras.layers.BatchNormalization(momentum=0.99, epsilon=0.001, name=f'{name}.tpu_batch_normalization')(x)
        x = tf.keras.layers.Lambda(lambda x: tf.nn.swish(x), name=f'{name}.swish')(x)
        return x
    return backbone_stem


def BackboneSE(reduce_filters, expand_filters, name=None):
    """EfficientNet Squeeze-Excitation block
    """
    def backbone_se(x_in):
        x = tf.keras.layers.Lambda(lambda x: tf.reduce_mean(x, [1, 2], keepdims=True))(x_in)
        x = tf.keras.layers.Conv2D(reduce_filters, 1, padding='same', use_bias=True, name=f'{name}.conv2d')(x)
        x = tf.keras.layers.Lambda(lambda x: tf.nn.swish(x))(x)
        x = tf.keras.layers.Conv2D(expand_filters, 1, padding='same', use_bias=True, name=f'{name}.conv2d_1')(x)
        x = tf.keras.layers.Lambda(lambda x: tf.nn.sigmoid(x))(x)
        x = tf.keras.layers.Multiply()([x, x_in])
        return x
    return backbone_se


def BackboneBlock(expand_filters, depth_ksize, se_filters, project_filters, strides, skip, name=None):
    """EfficientNet convolutional block
    """
    def backbone_block(x_in):
        x = x_in
        if expand_filters is not None:
            x = tf.keras.layers.Conv2D(expand_filters, 1, padding='same', use_bias=False, name=f'{name}.conv2d')(x)
            x = tf.keras.layers.BatchNormalization(momentum=0.99, epsilon=0.001, name=f'{name}.tpu_batch_normalization')(x)
        x = tf.keras.layers.DepthwiseConv2D(depth_ksize, strides, padding='same', use_bias=False, name=f'{name}.depthwise_conv2d')(x)
        x = tf.keras.layers.BatchNormalization(momentum=0.99, epsilon=0.001, name=f'{name}.tpu_batch_normalization{"_1" if expand_filters else ""}')(x)
        x = BackboneSE(*se_filters, name=f'{name}.se')(x)
        x = tf.keras.layers.Conv2D(project_filters, 1, padding='same', use_bias=False, name=f'{name}.conv2d{"_1" if expand_filters else ""}')(x)
        x = tf.keras.layers.BatchNormalization(momentum=0.99, epsilon=0.001, name=f'{name}.tpu_batch_normalization{"_2" if expand_filters else "_1"}')(x)
        if skip:
            x = tf.keras.layers.Add()([x_in, x])
        return x
    return backbone_block


def Backbone(stem_filters, repeats, expansions, se_filters, projections, name=None):
    """EfficientNet backbone model
    """
    def efficientnet(x_in):
        x = inputs = tf.keras.Input(x_in.shape[1:])
        x = BackboneStem(stem_filters, name='stem')(x)
        outputs = []
        strides = [1, 2, 2, 2, 1, 2, 1]
        depth_ksizes = [3, 3, 5, 3, 5, 5, 3]
        args = [repeats, projections, strides, depth_ksizes]
        counter = count(0)
        for i, (repeat, project, stride, depth_ksize) in enumerate(zip(*args)):
            for j in range(repeat):
                if i == 0:
                    expansion = None
                    se_filter = [se_filters[i], stem_filters] if j == 0 else [se_filters[i+1], project]
                else:
                    idx = i if j == 0 else i+1
                    expansion = expansions[idx-1]
                    se_filter = [se_filters[idx], expansions[idx-1]]
                stride = 1 if j != 0 else stride
                skip = j != 0 and stride == 1
                x = BackboneBlock(expansion, depth_ksize, se_filter, project, stride, skip, name=f'blocks_{next(counter)}')(x)
            if i in [2, 4, 6]:
                outputs.append(x)
        return tf.keras.Model(inputs=inputs, outputs=outputs, name=name)(x_in)
    return efficientnet


class ResampleFeatures(tf.keras.layers.Layer):
    """Resample layer
    """
    def __init__(self, filters, feature_level, name=None):
        super().__init__(name=name)
        self.filters = filters
        self.feature_level = feature_level

        self.conv_layer = tf.keras.layers.Conv2D(self.filters, 1, padding='same', name='conv2d')
        self.bn_layer = tf.keras.layers.BatchNormalization(momentum=0.99, epsilon=0.001, name='bn')
    
    def get_layer(self, name):
        for layer in [self.conv_layer, self.bn_layer]:
            if layer.name == name:
                return layer
        return None

    def call(self, x, features, training):
        height, width, channels = x.shape.as_list()[1:4]
        if not features:
            target_height, target_width = (height + 1) // 2, (width + 1) // 2
        else:
            target_height, target_width = features[self.feature_level].shape.as_list()[1:3]

        if channels != self.filters:
            x = self.conv_layer(x)
            x = self.bn_layer(x, training=training)
        if height > target_height and width > target_width:
            strides = (int((height - 1) // target_height + 1),
                       int((width - 1) // target_width + 1))
            x = tf.keras.layers.MaxPooling2D(pool_size=(strides[0]+1, strides[1]+1), strides=strides, padding='same')(x)
        elif height <= target_height and width <= target_width:
            x = tf.image.resize(x, (target_height, target_width), tf.image.ResizeMethod.NEAREST_NEIGHBOR)
        else:
            raise ValueError(f'Incompatible Resampling: {height}x{width}->{target_height}x{target_width}')
        return x


class FeaturePyramidNetwork(tf.keras.layers.Layer):
    """Feature pyramid network
    """
    def __init__(self, n_cells, filters, name=None):
        super().__init__(name=name)
        self.feature_levels = [3, 2, 1, 0, 1, 2, 3, 4]
        self.input_offsets = [[3, 4], [2, 5], [1, 6], [0, 7], [1, 7, 8], [2, 6, 9], [3, 5, 10], [4, 11]]
        self.cells = [FPNCell(filters, self.feature_levels, self.input_offsets, name=f'cell_{i}') for i in range(n_cells)]

    def get_layer(self, name):
        for cell in self.cells:
            if cell.name == name:
                return cell
        return None

    def call(self, x, training):
        for cell in self.cells:
            cell_x = cell(x, training)
            x = list()
            for level in range(3, 8):
                for i, node_level in enumerate(reversed(self.feature_levels)):
                    if node_level + 3 == level:
                        x.append(cell_x[-1 - i])
                        break
        return x


class FPNCell(tf.keras.layers.Layer):
    """Feature pyramid network cell
    """
    def __init__(self, filters, feature_levels, input_offsets, name=None):
        super().__init__(name=name)
        self.nodes = [FPNNode(filters, offset, level, name=f'fnode{i}') for i, (level, offset) in enumerate(zip(feature_levels, input_offsets))]

    def get_layer(self, name):
        for node in self.nodes:
            if node.name == name:
                return node
        return None

    def call(self, x, training):
        for node in self.nodes:
            x = node(x, training)
        return x


class FPNNode(tf.keras.layers.Layer):
    """Feature pyramid network node
    """
    def __init__(self, filters, input_offset, feature_level, name=None):
        super().__init__(name=name)
        self.filters = filters
        self.input_offset = input_offset
        self.feature_level = feature_level
        self.resample_layers = list()
        self.vars = list()
        self.conv = tf.keras.layers.SeparableConv2D(self.filters, 3, padding='same', name='conv')
        self.bn = tf.keras.layers.BatchNormalization(momentum=0.99, epsilon=0.001, name='bn')
    
    def get_layer(self, name):
        for layer in [*self.vars, *self.resample_layers, self.conv, self.bn]:
            if layer.name == name:
                return layer
        return None

    def build(self, input_shape):
        for i, offset in enumerate(self.input_offset):
            name = f'resample_{i}_{offset}_{len(input_shape)}'
            layer = ResampleFeatures(self.filters, self.feature_level, name)
            self.resample_layers.append(layer)
        for i in range(len(self.input_offset)):
            name = f'WSM{("" if i == 0 else f"_{i}")}'
            var = self.add_weight(shape=(), initializer='ones', name=name)
            self.vars.append(var)
        self.built = True
        super().build(input_shape)

    def fuse_inputs(self, x):
        dtype = x[0].dtype
        weights = [tf.nn.relu(tf.cast(var, dtype=dtype)) for var in self.vars]
        weights_sum = tf.add_n(weights)
        x = [x[i] * weights[i] / (weights_sum + 1e-4) for i in range(len(x))]
        x = tf.add_n(x)
        return x

    def conv_after_fuse(self, x, training):
        x = tf.nn.swish(x)
        x = self.conv(x)
        x = self.bn(x, training=training)
        return x

    def call(self, x, training):
        nodes = list()
        for i, offset in enumerate(self.input_offset):
            node = x[offset]
            node = self.resample_layers[i](node, x, training)
            nodes.append(node)
        node = self.fuse_inputs(nodes)
        node = self.conv_after_fuse(node, training)
        x = x + [node]
        return x


class ClassBlock(tf.keras.layers.Layer):
    """Classification head
    """
    def __init__(self, filters, repeats, classes, name=None):
        super().__init__(name=name)
        self.filters = filters
        self.repeats = repeats
        self.classes = classes
        self.n_anchors = 3
        self.conv_layers = list()
        self.bn_layers = list()
        for i in range(repeats):
            self.conv_layers.append(tf.keras.layers.SeparableConv2D(filters, 3, padding='same', name=f'class-{i}'))
            bn = list()
            for j in range(5):
                bn.append(tf.keras.layers.BatchNormalization(momentum=0.99, epsilon=0.001, name=f'class-{i}-bn-{j+3}'))
            self.bn_layers.append(bn)
        self.class_layer = tf.keras.layers.SeparableConv2D((classes+1)*self.n_anchors, 3, padding='same', name='class-predict')
    
    def get_layer(self, name):
        for layer in [*self.conv_layers, *[bn for bn_list in self.bn_layers for bn in bn_list], self.class_layer]:
            if layer.name == name:
                return layer
        return None

    def call(self, x_in, training):
        outputs = list()
        for level_id in range(5):
            x = x_in[level_id]
            for i in range(self.repeats):
                x = self.conv_layers[i](x)
                x = self.bn_layers[i][level_id](x, training=training)
                x = tf.nn.swish(x)
            x = self.class_layer(x)
            shape = tf.shape(x)
            x = tf.reshape(x, (-1, shape[1], shape[2], self.n_anchors, self.classes+1))
            outputs.append(x)
        return outputs


class RegressionBlock(tf.keras.layers.Layer):
    """Regression head
    """
    def __init__(self, filters, repeats, name=None):
        super().__init__(name=name)
        self.filters = filters
        self.repeats = repeats
        self.n_anchors = 3
        self.conv_layers = list()
        self.bn_layers = list()
        for i in range(repeats):
            self.conv_layers.append(tf.keras.layers.SeparableConv2D(filters, 3, padding='same', name=f'box-{i}'))
            bn = list()
            for j in range(5):
                bn.append(tf.keras.layers.BatchNormalization(momentum=0.99, epsilon=0.001, name=f'box-{i}-bn-{j+3}'))
            self.bn_layers.append(bn)
        self.regression_layer = tf.keras.layers.SeparableConv2D(4*self.n_anchors, 3, padding='same', name='box-predict')
    
    def get_layer(self, name):
        for layer in [*self.conv_layers, *[bn for bn_list in self.bn_layers for bn in bn_list], self.regression_layer]:
            if layer.name == name:
                return layer
        return None

    def call(self, x_in, training):
        outputs = list()
        for level_id in range(5):
            x = x_in[level_id]
            for i in range(self.repeats):
                x = self.conv_layers[i](x)
                x = self.bn_layers[i][level_id](x, training=training)
                x = tf.nn.swish(x)
            x = self.regression_layer(x)
            shape = tf.shape(x)
            x = tf.reshape(x, (-1, shape[1], shape[2], self.n_anchors, 4))
            outputs.append(x)
        return outputs


def Concat():
    """Concatenation layer
    """
    def concat(x_classes, x_boxes):
        x = list()
        for i, (x_c, x_b) in enumerate(zip(reversed(x_classes), reversed(x_boxes))):
            x_concat = tf.keras.layers.Concatenate(axis=-1, name=f'output_{i}')([x_b, x_c])
            x.append(x_concat)
        return x
    return concat


def EfficientDet(input_shape, classes, anchors, max_detections=512, iou_threshold=0.1, score_threshold=0.25, model='efficientdet-d0', training=False, postprocess=True):
    """EfficientDet object detector

    Args:
        input_shape (tuple): input tensor shape (K, L, C_in)
        classes (int): number of classes
        anchors (tf.Tensor): anchor boxes
        max_detections (int, optional): maximum output detections. Defaults to 512.
        iou_threshold (float, optional): non maximum suppression iou threshold. Defaults to 0.1.
        score_threshold (float, optional): minimum score of output detections. Defaults to 0.25.
        model (str, optional): model configuration, supports configurations D0 to D6, expects 'efficientdet-dX' where X is in range of 0-6. Defaults to 'efficientdet-d0'.
        training (bool, optional): training flag, if False - inference mode is set. Defaults to False.
        postprocess (bool, optional): wether to postprocess outputs during inference, if training is False and postprocess is False,
                                      outputs are equivalent to training mode outputs. Defaults to True.

    Returns:
        tf.keras.Model: EfficientDet model
    """
    models = {
        'efficientdet-d0': {
            'backbone': {
                'stem_filters': 32,
                'repeats': [1, 2, 2, 3, 3, 4, 1],
                'expansion': [96, 144, 240, 480, 672, 1152],
                'se_filters': [8, 4, 6, 10, 20, 28, 48],
                'projection': [16, 24, 40, 80, 112, 192, 320],
            },
            'resample_channels': 64,
            'feature_pyramide': {
                'fpn_cells': 3,
                'fpn_filters': 64,
            },
            'head': {
                'filters': 64,
                'repeats': 3,
            }
        },
        'efficientdet-d1': {
            'backbone': {
                'stem_filters': 32,
                'repeats': [2, 3, 3, 4, 4, 5, 2],
                'expansion': [96, 144, 240, 480, 672, 1152, 1920],
                'se_filters': [8, 4, 6, 10, 20, 28, 48, 80],
                'projection': [16, 24, 40, 80, 112, 192, 320],
            },
            'resample_channels': 88,
            'feature_pyramide': {
                'fpn_cells': 4,
                'fpn_filters': 88,
            },
            'head': {
                'filters': 88,
                'repeats': 3,
            }
        },
        'efficientdet-d2': {
            'backbone': {
                'stem_filters': 32,
                'repeats': [2, 3, 3, 4, 4, 5, 2],
                'expansion': [96, 144, 288, 528, 720, 1248, 2112],
                'se_filters': [8, 4, 6, 12, 22, 30, 52, 88],
                'projection': [16, 24, 48, 88, 120, 208, 352],
            },
            'resample_channels': 112,
            'feature_pyramide': {
                'fpn_cells': 5,
                'fpn_filters': 112,
            },
            'head': {
                'filters': 112,
                'repeats': 3,
            }
        },
        'efficientdet-d3': {
            'backbone': {
                'stem_filters': 40,
                'repeats': [2, 3, 3, 5, 5, 6, 2],
                'expansion': [144, 192, 288, 576, 816, 1392, 2304],
                'se_filters': [10, 6, 8, 12, 24, 34, 58, 96],
                'projection': [24, 32, 48, 96, 136, 232, 384],
            },
            'resample_channels': 160,
            'feature_pyramide': {
                'fpn_cells': 6,
                'fpn_filters': 160,
            },
            'head': {
                'filters': 160,
                'repeats': 4,
            }
        },
        'efficientdet-d4': {
            'backbone': {
                'stem_filters': 48,
                'repeats': [2, 4, 4, 6, 6, 8, 2],
                'expansion': [144, 192, 336, 672, 960, 1632, 2688],
                'se_filters': [12, 6, 8, 14, 28, 40, 68, 112],
                'projection': [24, 32, 56, 112, 160, 272, 448],
            },
            'resample_channels': 224,
            'feature_pyramide': {
                'fpn_cells': 7,
                'fpn_filters': 224,
            },
            'head': {
                'filters': 224,
                'repeats': 4,
            }
        },
        'efficientdet-d5': {
            'backbone': {
                'stem_filters': 48,
                'repeats': [3, 5, 5, 7, 7, 9, 3],
                'expansion': [144, 240, 384, 768, 1056, 1824, 3072],
                'se_filters': [12, 6, 10, 16, 32, 44, 76, 128],
                'projection': [24, 40, 64, 128, 176, 304, 512],
            },
            'resample_channels': 288,
            'feature_pyramide': {
                'fpn_cells': 7,
                'fpn_filters': 288,
            },
            'head': {
                'filters': 288,
                'repeats': 4,
            }
        },
        'efficientdet-d6': {
            'backbone': {
                'stem_filters': 56,
                'repeats': [3, 6, 6, 8, 8, 11, 3],
                'expansion': [192, 240, 432, 864, 1200, 2064, 3456],
                'se_filters': [14, 8, 10, 18, 36, 50, 86, 144],
                'projection': [32, 40, 72, 144, 200, 344, 576],
            },
            'resample_channels': 384,
            'feature_pyramide': {
                'fpn_cells': 8,
                'fpn_filters': 384,
            },
            'head': {
                'filters': 384,
                'repeats': 5,
            }
        },
    }
    x = inputs = tf.keras.layers.Input(input_shape, name=None)
    x = Backbone(*models[model]['backbone'].values(), name='efficientnet-b'+model[-1])(x)
    x_resample_1 = ResampleFeatures(models[model]['resample_channels'], 3, name='resample_p6')(x[-1], None, training)
    x_resample_2 = ResampleFeatures(models[model]['resample_channels'], 4, name='resample_p7')(x_resample_1, None, training)
    x = x + [x_resample_1, x_resample_2]
    x = FeaturePyramidNetwork(*models[model]['feature_pyramide'].values(), name='fpn_cells')(x, training)
    x_class = ClassBlock(*models[model]['head'].values(), classes, name='class_net')(x, training)
    x_regress = RegressionBlock(*models[model]['head'].values(), name='box_net')(x, training)
    outputs = Concat()(x_class, x_regress)
    if not training and postprocess:
        boxes_0 = tf.keras.layers.Lambda(lambda x: boxes_from_output(x, anchors[0, ...], classes), name='boxes_0')(outputs[0])
        boxes_1 = tf.keras.layers.Lambda(lambda x: boxes_from_output(x, anchors[1, ...], classes), name='boxes_1')(outputs[1])
        boxes_2 = tf.keras.layers.Lambda(lambda x: boxes_from_output(x, anchors[2, ...], classes), name='boxes_2')(outputs[2])
        boxes_3 = tf.keras.layers.Lambda(lambda x: boxes_from_output(x, anchors[3, ...], classes), name='boxes_3')(outputs[3])
        boxes_4 = tf.keras.layers.Lambda(lambda x: boxes_from_output(x, anchors[4, ...], classes), name='boxes_4')(outputs[4])
        boxes = [boxes_0, boxes_1, boxes_2, boxes_3, boxes_4]
        outputs = tf.keras.layers.Lambda(lambda x: non_maximum_suppression(x, classes, max_detections, iou_threshold, score_threshold), name='nms')(boxes)
    return tf.keras.Model(inputs=inputs, outputs=outputs, name=model)
