import tensorflow.keras.backend as K

from tensorflow.keras.models import Model

from tensorflow.keras import Input

from tensorflow.keras.layers import Conv2D, PReLU, UpSampling2D, concatenate , Reshape, Dense, Permute, MaxPool2D

from tensorflow.keras.layers import GlobalAveragePooling2D, Activation, add, GaussianNoise, BatchNormalization, multiply

from tensorflow.keras.optimizers import SGD

from loss import custom_loss

K.set_image_data_format("channels_last")

def unet_model(input_shape, modified_unet=True, learning_rate=0.01, start_channel=64, 

               number_of_levels=3, inc_rate=2, output_channels=4, saved_model_dir=None):

    """

    Builds UNet model

    Parameters

    ----------

    input_shape : tuple

        Shape of the input data (height, width, channel)

    modified_unet : bool

        Whether to use modified UNet or the original UNet

    learning_rate : float

        Learning rate for the model. The default is 0.01.

    start_channel : int

        Number of channels of the first conv. The default is 64.

    number_of_levels : int

        The depth size of the U-structure. The default is 3.

    inc_rate : int

        Rate at which the conv channels will increase. The default is 2.

    output_channels : int

        The number of output layer channels. The default is 4

    saved_model_dir : str

        If spesified, the model weights will be loaded from this path. The default is None.

    Returns

    -------

    model : keras.model

        The created keras model with respect to the input parameters

    """

    input_layer = Input(shape=input_shape, name='the_input_layer')

    if modified_unet:

        x = GaussianNoise(0.01, name='Gaussian_Noise')(input_layer)

        x = Conv2D(64, 2, padding='same')(x)

        x = level_block_modified(x, start_channel, number_of_levels, inc_rate)

        x = BatchNormalization(axis = -1)(x)

        x = PReLU(shared_axes=[1, 2])(x)

    else: 

        x = level_block(input_layer, start_channel, number_of_levels, inc_rate)

    x            = Conv2D(output_channels, 1, padding='same')(x)

    output_layer = Activation('softmax')(x)

    model        = Model(inputs = input_layer, outputs = output_layer)

    if modified_unet:

        print("The modified UNet was built!")

    else:

        print("The original UNet was built!")

    if saved_model_dir:

        model.load_weights(saved_model_dir)

        print("the model weights were successfully loaded!")

    sgd = SGD(lr=learning_rate, momentum=0.9, decay=0)

    model.compile(optimizer=sgd, loss=custom_loss)

    return model

def se_block(x, ratio=16):

    """

    creates a squeeze and excitation block

    https://arxiv.org/abs/1709.01507

    Parameters

    ----------

    x : tensor

        Input keras tensor

    ratio : int

        The reduction ratio. The default is 16.

    Returns

    -------

    x : tensor

        A keras tensor

    """

    channel_axis = 1 if K.image_data_format() == "channels_first" else -1

    filters = x.shape[channel_axis]

    se_shape = (1, 1, filters)

    se = GlobalAveragePooling2D()(x)

    se = Reshape(se_shape)(se)

    se = Dense(filters // ratio, activation='relu', kernel_initializer='he_normal', use_bias=False)(se)

    se = Dense(filters, activation='sigmoid', kernel_initializer='he_normal', use_bias=False)(se)

    if K.image_data_format() == 'channels_first':

        se = Permute((3, 1, 2))(se)

    x = multiply([x, se])

    return x

def level_block(x, dim, level, inc):

    if level > 0:

        m = conv_layers(x, dim)

        x = MaxPool2D(pool_size=(2, 2))(m)

        x = level_block(x,int(inc*dim), level-1, inc)

        x = UpSampling2D(size=(2, 2))(x)

        x = Conv2D(dim, 2, padding='same')(x)

        m = concatenate([m,x])

        x = conv_layers(m, dim)

    else:

        x = conv_layers(x, dim)

    return x

def level_block_modified(x, dim, level, inc):

    if level > 0:

        m = res_block(x, dim, encoder_path=True)##########

        x = Conv2D(int(inc*dim), 2, strides=2, padding='same')(m)

        x = level_block_modified(x, int(inc*dim), level-1, inc)

        x = UpSampling2D(size=(2, 2))(x)

        x = Conv2D(dim, 2, padding='same')(x)

        m = concatenate([m,x])

        m = se_block(m, 8)

        x = res_block(m, dim, encoder_path=False)

    else:

        x = res_block(x, dim, encoder_path=True) #############

    return x

def conv_layers(x, dim):

    x = Conv2D(dim, 3, padding='same')(x)

    x = Activation("relu")(x)

    x = Conv2D(dim, 3, padding='same')(x)

    x = Activation("relu")(x)

    return x

def res_block(x, dim, encoder_path=True):

    m = BatchNormalization(axis = -1)(x)

    m = PReLU(shared_axes = [1, 2])(m)

    m = Conv2D(dim, 3, padding='same')(m)

    m = BatchNormalization(axis = -1)(m)

    m = PReLU(shared_axes = [1, 2])(m)

    m = Conv2D(dim, 3, padding='same')(m)

    if encoder_path:

        x = add([x, m])

    else:

        x = Conv2D(dim, 1, padding='same', use_bias=False)(x)

        x = add([x,m])

    return  x