#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""YOLO_v3 VGG16 Model Defined in Keras."""

from tensorflow.keras.layers import Conv2D, UpSampling2D, Concatenate, MaxPooling2D
from tensorflow.keras.models import Model
from tensorflow.keras.applications.vgg16 import VGG16

from common.backbones.layers import YoloConv2D
from yolo3.models.layers import compose, DarknetConv2D, DarknetConv2D_BN_Leaky, make_last_layers


def yolo3_vgg16_body(inputs, num_anchors, num_classes):
    """Create YOLO_V3 model CNN body in Keras."""
    '''
    Layer Name input_1 Output: Tensor("input_1:0", shape=(?, 416, 416, 3), dtype=float32)
    Layer Name block1_conv1 Output: Tensor("block1_conv1/Relu:0", shape=(?, 416, 416, 64), dtype=float32)
    Layer Name block1_conv2 Output: Tensor("block1_conv2/Relu:0", shape=(?, 416, 416, 64), dtype=float32)
    Layer Name block1_pool Output: Tensor("block1_pool/MaxPool:0", shape=(?, 208, 208, 64), dtype=float32)
    Layer Name block2_conv1 Output: Tensor("block2_conv1/Relu:0", shape=(?, 208, 208, 128), dtype=float32)
    Layer Name block2_conv2 Output: Tensor("block2_conv2/Relu:0", shape=(?, 208, 208, 128), dtype=float32)
    Layer Name block2_pool Output: Tensor("block2_pool/MaxPool:0", shape=(?, 104, 104, 128), dtype=float32)
    Layer Name block3_conv1 Output: Tensor("block3_conv1/Relu:0", shape=(?, 104, 104, 256), dtype=float32)
    Layer Name block3_conv2 Output: Tensor("block3_conv2/Relu:0", shape=(?, 104, 104, 256), dtype=float32)
    Layer Name block3_conv3 Output: Tensor("block3_conv3/Relu:0", shape=(?, 104, 104, 256), dtype=float32)
    Layer Name block3_pool Output: Tensor("block3_pool/MaxPool:0", shape=(?, 52, 52, 256), dtype=float32)
    Layer Name block4_conv1 Output: Tensor("block4_conv1/Relu:0", shape=(?, 52, 52, 512), dtype=float32)
    Layer Name block4_conv2 Output: Tensor("block4_conv2/Relu:0", shape=(?, 52, 52, 512), dtype=float32)
    Layer Name block4_conv3 Output: Tensor("block4_conv3/Relu:0", shape=(?, 52, 52, 512), dtype=float32)
    Layer Name block4_pool Output: Tensor("block4_pool/MaxPool:0", shape=(?, 26, 26, 512), dtype=float32)
    Layer Name block5_conv1 Output: Tensor("block5_conv1/Relu:0", shape=(?, 26, 26, 512), dtype=float32)
    Layer Name block5_conv2 Output: Tensor("block5_conv2/Relu:0", shape=(?, 26, 26, 512), dtype=float32)
    Layer Name block5_conv3 Output: Tensor("block5_conv3/Relu:0", shape=(?, 26, 26, 512), dtype=float32)
    Layer Name block5_pool Output: Tensor("block5_pool/MaxPool:0", shape=(?, 13, 13, 512), dtype=float32)
    '''

    #net, endpoint = inception_v2.inception_v2(inputs)
    vgg16 = VGG16(input_tensor=inputs,weights='imagenet',include_top=False)
    x = vgg16.get_layer('block5_pool').output
    x = YoloConv2D(512, (3, 3), activation='relu', padding='same', name='block6_conv1')(x)
    x = YoloConv2D(512, (3, 3), activation='relu', padding='same', name='block6_conv2')(x)
    x = YoloConv2D(512, (3, 3), activation='relu', padding='same', name='block6_conv3')(x)
    x = YoloConv2D(512, (3, 3), activation='relu', padding='same', name='block6_conv4')(x)

    # input: 416 x 416 x 3
    # block6_conv3 :13 x 13 x 512
    # block5_conv3 :26 x 26 x 512
    # block4_conv3 : 52 x 52 x 512


    # f1 :13 x 13 x 1024 13 x 13 x 512
    x, y1 = make_last_layers(x, 512, num_anchors * (num_classes + 5), predict_id='1')

    x = compose(
            DarknetConv2D_BN_Leaky(256, (1,1)),
            UpSampling2D(2))(x)

    f2 = vgg16.get_layer('block5_conv3').output
    # f2: 26 x 26 x 512
    x = Concatenate()([x,f2])

    x, y2 = make_last_layers(x, 256, num_anchors*(num_classes+5), predict_id='2')

    x = compose(
            DarknetConv2D_BN_Leaky(128, (1,1)),
            UpSampling2D(2))(x)

    f3 = vgg16.get_layer('block4_conv3').output
    # f3 : 52 x 52 x 256
    x = Concatenate()([x, f3])
    x, y3 = make_last_layers(x, 128, num_anchors*(num_classes+5), predict_id='3')

    return Model(inputs = inputs, outputs=[y1,y2,y3])

def tiny_yolo3_vgg16_body(inputs, num_anchors, num_classes):
    '''Create Tiny YOLO_v3 VGG16 model CNN body in keras.'''
    vgg16 = VGG16(input_tensor=inputs,weights='imagenet',include_top=False)
    x = vgg16.get_layer('block5_pool').output
    x = YoloConv2D(512, (3, 3), activation='relu', padding='same', name='block6_conv1')(x)
    x = YoloConv2D(512, (3, 3), activation='relu', padding='same', name='block6_conv2')(x)
    x = YoloConv2D(512, (3, 3), activation='relu', padding='same', name='block6_conv3')(x)
    #x = YoloConv2D(512, (3, 3), activation='relu', padding='same', name='block6_conv4')(x)

    # input: 416 x 416 x 3
    # block6_conv3 :13 x 13 x 512
    # block5_conv3 :26 x 26 x 512
    # block4_conv3 : 52 x 52 x 512

    x1 = vgg16.get_layer('block5_conv3').output
    x2 = x

    x2 = DarknetConv2D_BN_Leaky(512, (1,1))(x2)

    y1 = compose(
            DarknetConv2D_BN_Leaky(1024, (3,3)),
            #Depthwise_Separable_Conv2D_BN_Leaky(filters=1024, kernel_size=(3, 3), block_id_str='14'),
            DarknetConv2D(num_anchors*(num_classes+5), (1,1)))(x2)

    x2 = compose(
            DarknetConv2D_BN_Leaky(256, (1,1)),
            UpSampling2D(2))(x2)
    y2 = compose(
            Concatenate(),
            DarknetConv2D_BN_Leaky(512, (3,3)),
            #Depthwise_Separable_Conv2D_BN_Leaky(filters=512, kernel_size=(3, 3), block_id_str='15'),
            DarknetConv2D(num_anchors*(num_classes+5), (1,1)))([x2,x1])

    return Model(inputs, [y1,y2])


tensorflow.keras.layers.Concatenate tensorflow.keras.layers.UpSampling2D tensorflow.keras.applications.vgg16.VGG16 tensorflow.keras.models.Model
[(38, 'tensorflow.keras.applications.vgg16.VGG16', 'VGG16', ([], {'input_tensor': 'inputs', 'weights': '"""imagenet"""', 'include_top': '(False)'}), False, 'from tensorflow.keras.applications.vgg16 import VGG16\n'), (52, 'yolo3.models.layers.make_last_layers', 'make_last_layers', (['x', '(512)', '(num_anchors * (num_classes + 5))'], {'predict_id': '"""1"""'}), False, 'from yolo3.models.layers import compose, DarknetConv2D, DarknetConv2D_BN_Leaky, make_last_layers\n'), (62, 'yolo3.models.layers.make_last_layers', 'make_last_layers', (['x', '(256)', '(num_anchors * (num_classes + 5))'], {'predict_id': '"""2"""'}), False, 'from yolo3.models.layers import compose, DarknetConv2D, DarknetConv2D_BN_Leaky, make_last_layers\n'), (71, 'yolo3.models.layers.make_last_layers', 'make_last_layers', (['x', '(128)', '(num_anchors * (num_classes + 5))'], {'predict_id': '"""3"""'}), False, 'from yolo3.models.layers import compose, DarknetConv2D, DarknetConv2D_BN_Leaky, make_last_layers\n'), (73, 'tensorflow.keras.models.Model', 'Model', ([], {'inputs': 'inputs', 'outputs': '[y1, y2, y3]'}), False, 'from tensorflow.keras.models import Model\n'), (77, 'tensorflow.keras.applications.vgg16.VGG16', 'VGG16', ([], {'input_tensor': 'inputs', 'weights': '"""imagenet"""', 'include_top': '(False)'}), False, 'from tensorflow.keras.applications.vgg16 import VGG16\n'), (108, 'tensorflow.keras.models.Model', 'Model', (['inputs', '[y1, y2]'], {}), False, 'from tensorflow.keras.models import Model\n'), (40, 'common.backbones.layers.YoloConv2D', 'YoloConv2D', (['(512)', '(3, 3)'], {'activation': '"""relu"""', 'padding': '"""same"""', 'name': '"""block6_conv1"""'}), False, 'from common.backbones.layers import YoloConv2D\n'), (41, 'common.backbones.layers.YoloConv2D', 'YoloConv2D', (['(512)', '(3, 3)'], {'activation': '"""relu"""', 'padding': '"""same"""', 'name': '"""block6_conv2"""'}), False, 'from common.backbones.layers import YoloConv2D\n'), (42, 'common.backbones.layers.YoloConv2D', 'YoloConv2D', (['(512)', '(3, 3)'], {'activation': '"""relu"""', 'padding': '"""same"""', 'name': '"""block6_conv3"""'}), False, 'from common.backbones.layers import YoloConv2D\n'), (43, 'common.backbones.layers.YoloConv2D', 'YoloConv2D', (['(512)', '(3, 3)'], {'activation': '"""relu"""', 'padding': '"""same"""', 'name': '"""block6_conv4"""'}), False, 'from common.backbones.layers import YoloConv2D\n'), (60, 'tensorflow.keras.layers.Concatenate', 'Concatenate', ([], {}), False, 'from tensorflow.keras.layers import Conv2D, UpSampling2D, Concatenate, MaxPooling2D\n'), (70, 'tensorflow.keras.layers.Concatenate', 'Concatenate', ([], {}), False, 'from tensorflow.keras.layers import Conv2D, UpSampling2D, Concatenate, MaxPooling2D\n'), (79, 'common.backbones.layers.YoloConv2D', 'YoloConv2D', (['(512)', '(3, 3)'], {'activation': '"""relu"""', 'padding': '"""same"""', 'name': '"""block6_conv1"""'}), False, 'from common.backbones.layers import YoloConv2D\n'), (80, 'common.backbones.layers.YoloConv2D', 'YoloConv2D', (['(512)', '(3, 3)'], {'activation': '"""relu"""', 'padding': '"""same"""', 'name': '"""block6_conv2"""'}), False, 'from common.backbones.layers import YoloConv2D\n'), (81, 'common.backbones.layers.YoloConv2D', 'YoloConv2D', (['(512)', '(3, 3)'], {'activation': '"""relu"""', 'padding': '"""same"""', 'name': '"""block6_conv3"""'}), False, 'from common.backbones.layers import YoloConv2D\n'), (92, 'yolo3.models.layers.DarknetConv2D_BN_Leaky', 'DarknetConv2D_BN_Leaky', (['(512)', '(1, 1)'], {}), False, 'from yolo3.models.layers import compose, DarknetConv2D, DarknetConv2D_BN_Leaky, make_last_layers\n'), (55, 'yolo3.models.layers.DarknetConv2D_BN_Leaky', 'DarknetConv2D_BN_Leaky', (['(256)', '(1, 1)'], {}), False, 'from yolo3.models.layers import compose, DarknetConv2D, DarknetConv2D_BN_Leaky, make_last_layers\n'), (56, 'tensorflow.keras.layers.UpSampling2D', 'UpSampling2D', (['(2)'], {}), False, 'from tensorflow.keras.layers import Conv2D, UpSampling2D, Concatenate, MaxPooling2D\n'), (65, 'yolo3.models.layers.DarknetConv2D_BN_Leaky', 'DarknetConv2D_BN_Leaky', (['(128)', '(1, 1)'], {}), False, 'from yolo3.models.layers import compose, DarknetConv2D, DarknetConv2D_BN_Leaky, make_last_layers\n'), (66, 'tensorflow.keras.layers.UpSampling2D', 'UpSampling2D', (['(2)'], {}), False, 'from tensorflow.keras.layers import Conv2D, UpSampling2D, Concatenate, MaxPooling2D\n'), (95, 'yolo3.models.layers.DarknetConv2D_BN_Leaky', 'DarknetConv2D_BN_Leaky', (['(1024)', '(3, 3)'], {}), False, 'from yolo3.models.layers import compose, DarknetConv2D, DarknetConv2D_BN_Leaky, make_last_layers\n'), (97, 'yolo3.models.layers.DarknetConv2D', 'DarknetConv2D', (['(num_anchors * (num_classes + 5))', '(1, 1)'], {}), False, 'from yolo3.models.layers import compose, DarknetConv2D, DarknetConv2D_BN_Leaky, make_last_layers\n'), (100, 'yolo3.models.layers.DarknetConv2D_BN_Leaky', 'DarknetConv2D_BN_Leaky', (['(256)', '(1, 1)'], {}), False, 'from yolo3.models.layers import compose, DarknetConv2D, DarknetConv2D_BN_Leaky, make_last_layers\n'), (101, 'tensorflow.keras.layers.UpSampling2D', 'UpSampling2D', (['(2)'], {}), False, 'from tensorflow.keras.layers import Conv2D, UpSampling2D, Concatenate, MaxPooling2D\n'), (103, 'tensorflow.keras.layers.Concatenate', 'Concatenate', ([], {}), False, 'from tensorflow.keras.layers import Conv2D, UpSampling2D, Concatenate, MaxPooling2D\n'), (104, 'yolo3.models.layers.DarknetConv2D_BN_Leaky', 'DarknetConv2D_BN_Leaky', (['(512)', '(3, 3)'], {}), False, 'from yolo3.models.layers import compose, DarknetConv2D, DarknetConv2D_BN_Leaky, make_last_layers\n'), (106, 'yolo3.models.layers.DarknetConv2D', 'DarknetConv2D', (['(num_anchors * (num_classes + 5))', '(1, 1)'], {}), False, 'from yolo3.models.layers import compose, DarknetConv2D, DarknetConv2D_BN_Leaky, make_last_layers\n')]