from tensorflow.keras import backend as K

from tensorflow.keras.applications.densenet import DenseNet121

from tensorflow.keras.applications.densenet import DenseNet169

# Model Imports

from tensorflow.keras.applications.resnet_v2 import ResNet50V2

from tensorflow.keras.layers import Dense, GlobalAveragePooling2D, Multiply, Concatenate, Input

from tensorflow.keras.models import Model

# Parallelize Custom Block

def parallelize_block(tensor, ratio=128):

    nb_channel = K.int_shape(tensor)[-1]

    # Averaging

    x = GlobalAveragePooling2D()(tensor)

    x = Dense(nb_channel // ratio, activation='relu')(x)

    x = Dense(nb_channel, activation='sigmoid')(x)

    # Weighting

    x = Multiply()([tensor, x])

    return x

# MAIN MODEL

def ParallelXNet(n_classes, tau=128, input_shape=(320, 320, 3), weights='imagenet', output_activation="sigmoid"):

    # Combined Input

    combinedInput = Input(batch_shape=(None, None, None, 3))

    # Base models preoaded with imagenet weights

    resnetV2_m = ResNet50V2(weights=weights, include_top=False,

                            input_shape=input_shape, input_tensor=combinedInput)

    densenet121_m = DenseNet121(weights=weights, include_top=False,

                                input_shape=input_shape, input_tensor=combinedInput)

    densenet169_m = DenseNet169(weights=weights, include_top=False,

                                input_shape=input_shape, input_tensor=combinedInput)

    for layer in resnetV2_m.layers:

        layer._name = layer.name + str("_0")

    for layer in densenet121_m.layers:

        layer._name = layer.name + str("_1")

    for layer in densenet169_m.layers:

        layer._name = layer.name + str("_2")

    base_model_output = [

        resnetV2_m.output,

        densenet121_m.output,

        densenet169_m.output

    ]

    # Concatenating the models

    c_model = Concatenate()(base_model_output)

    # Passing via parallelize block

    par_c_model = parallelize_block(

        tensor=c_model,

        ratio=tau

    )

    # Final Average Pooling layer

    x = GlobalAveragePooling2D()(par_c_model)

    # Final Dense Layer

    detection_output = Dense(n_classes, activation=output_activation)(x)

    return Model(inputs=combinedInput, outputs=detection_output)