import tensorflow as tf

from itertools import count

from .detector_utils import boxes_from_output, non_maximum_suppression


def spatial2depth(x):
    """Concentrates spacial information in channel dimension: (b, w, h, c) -> (b, w/2, h/2, 4c)

    Args:
        x (tf.Tensor): input tensor

    Returns:
        tf.Tensor: output tensor
    """
    # (b, w, h, c) -> (b, w/2, h/2, 4c)
    x1 = x[:, ::2, ::2, :]
    x2 = x[:, 1::2, ::2, :]
    x3 = x[:, ::2, 1::2, :]
    x4 = x[:, 1::2, 1::2, :]
    return tf.concat([x1, x2, x3, x4], axis=-1)


def ConvBlock(filters, ksize, strides=1, groups=1, name=None):
    """YOLOv5 Convolutional block
    """
    def conv_block(x):
        x = tf.keras.layers.ZeroPadding2D(ksize // 2, name=f'{name}.padding')(x)
        x = tf.keras.layers.Conv2D(filters, ksize, strides=strides, padding='valid', groups=groups, use_bias=False, name=f'{name}.conv')(x)
        x = tf.keras.layers.BatchNormalization(epsilon=1e-3, momentum=0.03, name=f'{name}.bn')(x)
        x = tf.keras.layers.Lambda(lambda x: tf.nn.silu(x), name=f'{name}.silu')(x)
        return x
    return conv_block


def BottleneckBlock(filters, groups=1, expansion=0.5, shortcut=True, name=None):
    """YOLOv5 Bottleneck convolutional block
    """
    def bottleneck_block(x_in):
        x = ConvBlock(filters * expansion, 1, name=f'{name}.cv1')(x_in)
        x = ConvBlock(filters, 3, groups=groups, name=f'{name}.cv2')(x)
        if shortcut:
            x = tf.keras.layers.Add(name=f'{name}.add')([x_in, x])
        return x
    return bottleneck_block


def C3Block(filters, n=1, groups=1, expansion=0.5, shortcut=True, name=None):
    """YOLOv5 C3 convolutional block
    """
    def c3_block(x_in):
        x1 = ConvBlock(filters * expansion, 1, name=f'{name}.cv1')(x_in)
        for i in range(n):
            x1 = BottleneckBlock(filters * expansion, groups=groups, expansion=1, shortcut=shortcut, name=f'{name}.m.{i}')(x1)
        x2 = ConvBlock(filters * expansion, 1, name=f'{name}.cv2')(x_in)
        x = tf.keras.layers.Concatenate(axis=-1, name=f'{name}.concat')([x1, x2])
        x = ConvBlock(filters, 1, name=f'{name}.cv3')(x)
        return x
    return c3_block


def SpatialPyramidePooling(filters, name=None):
    """YOLOv5 Pooling block
    """
    def spp(x):
        x = ConvBlock(filters // 2, 1, name=f'{name}.cv1')(x)
        ksizes = (5, 9, 13)
        x_list = [x]
        for ksize in ksizes:
            xk = tf.keras.layers.ZeroPadding2D(ksize // 2, name=f'{name}.{ksize}.padding')(x)
            xk = tf.keras.layers.MaxPool2D(ksize, strides=1, name=f'{name}.{ksize}.maxpool')(xk)
            x_list.append(xk)
        x = tf.keras.layers.Concatenate(axis=-1, name=f'{name}.concat')(x_list)
        x = ConvBlock(filters, 1, name=f'{name}.cv2')(x)
        return x
    return spp


def Backbone(filters, n, name=None):
    """YOLOv5 Backbone
    """
    def yolo_backbone(x_in):
        nonlocal filters
        x = inputs = tf.keras.Input(x_in.shape[1:])
        # focus
        x = tf.keras.layers.Lambda(lambda x: spatial2depth(x), name='yolov5.0.spatial2depth')(x)
        x = ConvBlock(filters, 3, name='yolov5.0.conv')(x)
        # 3x conv, c3
        filters *= 2
        x = ConvBlock(filters, 3, strides=2, name='yolov5.1')(x)
        x = C3Block(filters, n=n, name='yolov5.2')(x)
        filters *= 2
        x = ConvBlock(filters, 3, strides=2, name='yolov5.3')(x)
        x = x4 = C3Block(filters, n=3*n, name='yolov5.4')(x)
        filters *= 2
        x = ConvBlock(filters, 3, strides=2, name='yolov5.5')(x)
        x = x6 = C3Block(filters, n=3*n, name='yolov5.6')(x)
        # conv, spp, c3
        filters *= 2
        x = ConvBlock(filters, 3, strides=2, name='yolov5.7')(x)
        x = SpatialPyramidePooling(filters, name='yolov5.8')(x)
        x = x9 = C3Block(filters, n=n, shortcut=False, name='yolov5.9')(x)
        return tf.keras.Model(inputs=inputs, outputs=[x4, x6, x9], name=name)(x_in)
    return yolo_backbone


def YoloBlock(filters, n, layer_id, upsample=True, name=None):
    """YOLOv5 Main block
    """
    def yolo_block(x_in):
        i = count(layer_id)
        x_skip, x = inputs = tf.keras.Input(x_in[0].shape[1:]), tf.keras.Input(x_in[1].shape[1:])
        x = x_prev = ConvBlock(filters, 1 if upsample else 3, 1 if upsample else 2, name=f'yolov5.{next(i)}')(x)
        if upsample:
            x = tf.keras.layers.UpSampling2D(2, interpolation='nearest', name=f'yolov5.{next(i)}')(x)
        x = tf.keras.layers.Concatenate(axis=-1, name=f'yolov5.{next(i)}')([x, x_skip])
        x = C3Block(filters if upsample else 2*filters, n=n, shortcut=False, name=f'yolov5.{next(i)}')(x)
        return tf.keras.Model(inputs=inputs, outputs=[x_prev, x] if upsample else x, name=name)(x_in)
    return yolo_block


def YoloOutput(classes, layer_id, name=None):
    """YOLOv5 Output block
    """
    def yolo_output(x_in):
        x = inputs = tf.keras.Input(x_in.shape[1:])
        x = tf.keras.layers.Conv2D((classes + 5) * 3, 1, name=f'yolov5.24.m.{layer_id}')(x)
        x = tf.keras.layers.Lambda(lambda x: tf.reshape(x, (-1, tf.shape(x)[1], tf.shape(x)[2], 3, classes + 5)), name=f'yolov5.output_{layer_id}.reshape')(x)
        return tf.keras.Model(inputs=inputs, outputs=x, name=name)(x_in)
    return yolo_output


def YOLOv5(input_shape, classes, anchors=None, max_detections=512, iou_threshold=0.1, score_threshold=0.25, model='yolov5s', training=False, postprocess=True):
    """YOLOv5 Object detector model

    Args:
        input_shape (tuple): model's input shape (K, L, C_in)
        classes (int): number of classes
        anchors (tf.Tensor, optional): anchor boxes. Defaults to None.
        max_detections (int, optional): maximum output detections. Defaults to 512.
        iou_threshold (float, optional): non max suppression iou threshold. Defaults to 0.1.
        score_threshold (float, optional): minimum object score to output. Defaults to 0.25.
        model (str, optional): model configuration, 'yolov5-s', 'yolov5-m' or 'yolov5-l'. Defaults to 'yolov5s'.
        training (bool, optional): training or testing mode flag, True equals to training mode, otherwise testing mode. Defaults to False.
        postprocess (bool, optional): flag, wether to posprocess outputs during testing, postprocess=True and
                                      training=False yields equivalent results to training=True. Defaults to True.

    Returns:
        _type_: _description_
    """
    models = {
        'yolov5s': {'filters': 32, 'n': 1},
        'yolov5m': {'filters': 48, 'n': 2},
        'yolov5l': {'filters': 64, 'n': 3}
    }
    filters, n = models[model]['filters'], models[model]['n']
    x = inputs = tf.keras.layers.Input(input_shape, name='yolo_input')
    x4, x6, x = Backbone(filters, n, name='yolo_backbone')(x)
    x10, x = YoloBlock(filters * 8, n, 10, upsample=True, name='yolo_conv_0')([x6, x])
    x14, x17 = YoloBlock(filters * 4, n, 14, upsample=True, name='yolo_conv_1')([x4, x])
    x = x20 = YoloBlock(filters * 4, n, 18, upsample=False, name='yolo_conv_2')([x14, x17])
    x = x23 = YoloBlock(filters * 8, n, 21, upsample=False, name='yolo_conv_3')([x10, x])
    output_0 = YoloOutput(classes, 0, name='yolo_output_0')(x17)
    output_1 = YoloOutput(classes, 1, name='yolo_output_1')(x20)
    output_2 = YoloOutput(classes, 2, name='yolo_output_2')(x23)
    if training or not postprocess:
        outputs = [output_2, output_1, output_0]
    else:
        boxes_0 = tf.keras.layers.Lambda(lambda x: boxes_from_output(x, anchors[2, ...], classes), name='yolo_boxes_0')(output_0)
        boxes_1 = tf.keras.layers.Lambda(lambda x: boxes_from_output(x, anchors[1, ...], classes), name='yolo_boxes_1')(output_1)
        boxes_2 = tf.keras.layers.Lambda(lambda x: boxes_from_output(x, anchors[0, ...], classes), name='yolo_boxes_2')(output_2)
        outputs = tf.keras.layers.Lambda(lambda x: non_maximum_suppression(x, classes, max_detections, iou_threshold, score_threshold), name='yolo_nms')([boxes_2, boxes_1, boxes_0])
    return tf.keras.Model(inputs=inputs, outputs=outputs, name=model)
