from keras import Model, Input
from keras.layers import Concatenate
from keras.layers import Conv2D, Reshape, TimeDistributed, Flatten, \
    Dense
from tensorflow import Tensor
from tensorflow.python import RandomNormal

from nets.classifier import RoiPoolingConv
from nets.vgg import VGG16_if as backbones, vgg_classifier_layers as classifier_layers


def IF_RPN(baselayers, num_anchors):
    """
    :param num_anchors: int, num box per anchor
    :param baselayers: List of feature maps [x0,x1]
    :return: rpn out
    :feature format: channel last
    """
    # get feature map
    assert len(baselayers) == 2, 'output of backbones must be list of 2 feature maps'
    x0 = baselayers[0]
    x1 = baselayers[1]
    print(x0.shape, x1.shape)
    """
    :param rpn_out : list [regression, classes]
    """
    rpn_out_0 = get_rpn_fusion(x0, num_anchors=num_anchors, count=0)
    rpn_out_1 = get_rpn_fusion(x1, num_anchors=num_anchors, count=1)
    """
    :returns fusion rpn output :List [regression, classes]
    """
    regression = Concatenate(axis=-2, name='classification')([rpn_out_0[0], rpn_out_1[0]])
    classes = Concatenate(axis=-2, name='regression')([rpn_out_0[1], rpn_out_1[1]])

    return [regression, classes]


def get_if_rcnn(inputs: Tensor):
    """
    :param inputs: Tensor from Input Layer
    :return:
    """
    # get back bone outputs
    if_backbones_out = backbones(inputs)

    return if_backbones_out


def get_rpn_fusion(base_layers, num_anchors, count=0):
    # ----------------------------------------------------#
    #   利用一个512通道的3x3卷积进行特征整合
    # ----------------------------------------------------#
    x2 = Conv2D(256,
                kernel_size=(3, 3),
                padding='same',
                activation='relu',
                kernel_initializer=RandomNormal(stddev=0.02)
                , name='rpn_conv1_' + str(count) + '_fusion2'
                )(base_layers)
    x3 = Conv2D(256, (5, 5), padding='same',
                activation='relu',
                kernel_initializer=RandomNormal(stddev=0.02)
                , name='rpn_conv1' + str(count) + '_fusion3'
                )(base_layers)
    x1 = Conv2D(256, (1, 1),
                padding='same',
                activation='relu',
                kernel_initializer=RandomNormal(stddev=0.02)
                , name='rpn_conv1' + str(count) + '_fusion1'
                )(base_layers)
    x = Concatenate(axis=-1)([x1, x2, x3])
    # ----------------------------------------------------#
    #   利用一个1x1卷积调整通道数，获得预测结果
    # ----------------------------------------------------#
    x_class = Conv2D(num_anchors, (1, 1), activation='sigmoid', kernel_initializer=RandomNormal(stddev=0.02)
                     , name='rpn_out_class_' + str(count)
                     )(x)
    x_regr = Conv2D(num_anchors * 4, (1, 1), activation='linear', kernel_initializer=RandomNormal(stddev=0.02)
                    , name='rpn_out_regress_' + str(count)
                    )(x)

    x_class = Reshape((-1, 1)
                      , name="classification_" + str(count)
                      )(x_class)
    x_regr = Reshape((-1, 4)
                     , name="regression_" + str(count)
                     )(x_regr)
    return [x_class, x_regr]


def get_if_classifier_vgg(base_layers, input_rois, roi_size, num_classes=21):
    # batch_size, 38, 38, 1024 -> batch_size, num_rois, 14, 14, 1024
    out_roi_pool0 = RoiPoolingConv(roi_size)([base_layers[0], input_rois])
    out_roi_pool1 = RoiPoolingConv(roi_size)([base_layers[1], input_rois])
    out_roi_pool = Concatenate(axis=1)([out_roi_pool0, out_roi_pool1])

    # batch_size, num_rois, 14, 14, 1024 -> num_rois, 1, 1, 2048
    out = classifier_layers(out_roi_pool)

    # batch_size, num_rois, 1, 1, 2048 -> batch_size, num_rois, 2048
    out = TimeDistributed(Flatten())(out)

    # batch_size, num_rois, 2048 -> batch_size, num_rois, num_classes
    out_class = TimeDistributed(Dense(num_classes, activation='softmax', kernel_initializer=RandomNormal(stddev=0.02)),
                                name='dense_class_{}'.format(num_classes))(out)
    # batch_size, num_rois, 2048 -> batch_size, num_rois, 4 * (num_classes-1)
    out_regr = TimeDistributed(
        Dense(4 * (num_classes - 1), activation='linear', kernel_initializer=RandomNormal(stddev=0.02)),
        name='dense_regress_{}'.format(num_classes))(out)
    return [out_class, out_regr]


def get_model_if_vgg(num_classes, num_anchors=6):
    inputs = Input(shape=(None, None, 3))
    inputs = Input(shape=(640, 640, 3))

    roi_input = Input(shape=(None, 4))

    # ----------------------------------------------------#
    #   假设输入为600,600,3
    #   获得一个38,38,1024的共享特征层base_layers
    # ----------------------------------------------------#
    base_layers = backbones(inputs)
    # print(base_layers[0].shape,base_layers[1].shape)
    # ----------------------------------------------------#
    #   将共享特征层传入建议框网络
    #   该网络结果会对先验框进行调整获得建议框
    # ----------------------------------------------------#
    rpn = IF_RPN(base_layers, num_anchors)
    # ----------------------------------------------------#
    #   将共享特征层和建议框传入classifier网络
    #   该网络结果会对建议框进行调整获得预测框
    # ----------------------------------------------------#
    classifier = get_if_classifier_vgg(base_layers, roi_input, 14, num_classes)

    model_rpn = Model(inputs, rpn)
    model_all = Model([inputs, roi_input], rpn + classifier)
    return model_rpn, model_all


if __name__ == '__main__':
    model_rpn, model_all = get_model_if_vgg(4)

    model_all.summary()
