"""
Created on 2021/3/15 9:58.
@Author: haifei
"""


from tensorflow.keras.layers import Input, Conv2D, BatchNormalization, Activation
from tensorflow.keras.layers import UpSampling2D, ZeroPadding2D, DepthwiseConv2D
from tensorflow.keras import Model


def _conv_block(inputs, filters, alpha, kernel=(3, 3), strides=(1, 1)):
    filters = int(filters * alpha)

    # conv + bn + relu <-- yolo v3 darknet
    x = ZeroPadding2D(padding=(1, 1))(inputs)
    x = Conv2D(filters, kernel, padding='valid', strides=strides)(x)
    x = BatchNormalization()(x)
    x = Activation("relu")(x)

    return x


def _depthwise_conv_block(inputs, pointwise_conv_filters, alpha, depth_multiplier=1, strides=(1, 1)):
    pointwise_conv_filters = int(pointwise_conv_filters * alpha)

    x = ZeroPadding2D((1, 1))(inputs)
    x = DepthwiseConv2D((3, 3), padding='valid', depth_multiplier=depth_multiplier, strides=strides)(x)
    x = BatchNormalization()(x)
    x = Activation("relu")(x)

    x = Conv2D(pointwise_conv_filters, (1, 1), padding='same', strides=(1, 1))(x)
    x = BatchNormalization()(x)
    x = Activation("relu")(x)

    return x


def get_mobilenet_encoder(inputs):  # mobilenetv1=mobilenets
    alpha = 1.0
    depth_multiplier = 1

    x = _conv_block(inputs, 32, alpha, strides=(2, 2))  # 下采样
    x = _depthwise_conv_block(x, 64, alpha, depth_multiplier)
    f1 = x

    x = _depthwise_conv_block(x, 128, alpha, depth_multiplier, strides=(2, 2))  # 下采样
    x = _depthwise_conv_block(x, 128, alpha, depth_multiplier)
    f2 = x

    x = _depthwise_conv_block(x, 256, alpha, depth_multiplier, strides=(2, 2))  # 下采样
    x = _depthwise_conv_block(x, 256, alpha, depth_multiplier)
    f3 = x

    x = _depthwise_conv_block(x, 512, alpha, depth_multiplier, strides=(2, 2))  # 下采样
    x = _depthwise_conv_block(x, 512, alpha, depth_multiplier)
    x = _depthwise_conv_block(x, 512, alpha, depth_multiplier)
    x = _depthwise_conv_block(x, 512, alpha, depth_multiplier)
    x = _depthwise_conv_block(x, 512, alpha, depth_multiplier)
    x = _depthwise_conv_block(x, 512, alpha, depth_multiplier)
    f4 = x

    x = _depthwise_conv_block(x, 1024, alpha, depth_multiplier, strides=(2, 2))
    x = _depthwise_conv_block(x, 1024, alpha, depth_multiplier)
    f5 = x

    return [f1, f2, f3, f4, f5]



def get_segnet_decoder(feature):
    #
    x = Conv2D(512, (3, 3), strides=1, padding='same', activation='relu')(feature)
    x = BatchNormalization()(x)
    #
    x = UpSampling2D(size=(2, 2))(x)
    x = Conv2D(256, (3, 3), strides=1, padding='same', activation='relu')(x)
    x = BatchNormalization()(x)
    #
    x = UpSampling2D(size=(2, 2))(x)
    x = Conv2D(128, (3, 3), strides=1, padding='same', activation='relu')(x)
    x = BatchNormalization()(x)
    #
    x = UpSampling2D(size=(2, 2))(x)
    x = Conv2D(128, (3, 3), strides=1, padding='same', activation='relu')(x)
    x = BatchNormalization()(x)
    #
    x = UpSampling2D(size=(2, 2))(x)
    x = Conv2D(64, (3, 3), strides=1, padding='same', activation='relu')(x)
    x = BatchNormalization()(x)
    #
    return x


def build_model(tif_size, bands, class_num):
    from pathlib import Path
    import sys
    print('===== %s =====' % Path(__file__).name)
    print('===== %s =====' % sys._getframe().f_code.co_name)

    # 输入
    inputs = Input(shape=(tif_size, tif_size, bands))
    # 编码器
    levels = get_mobilenet_encoder(inputs)
    # 解码器
    x = get_segnet_decoder(feature=levels[3])
    # 输出
    x = Conv2D(class_num, (1, 1), strides=1, padding='same', activation='softmax')(x)

    mymodel = Model(inputs, x)
    return mymodel