""" 

Copyright (C) 2022 King Saud University, Saudi Arabia 

SPDX-License-Identifier: Apache-2.0 

Licensed under the Apache License, Version 2.0 (the "License"); you may not use

this file except in compliance with the License. You may obtain a copy of the 

License at

http://www.apache.org/licenses/LICENSE-2.0  

Unless required by applicable law or agreed to in writing, software distributed

under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR 

CONDITIONS OF ANY KIND, either express or implied. See the License for the

specific language governing permissions and limitations under the License. 

Author:  Hamdi Altaheri 

"""

#%%

import math

import tensorflow as tf

from tensorflow.keras.layers import GlobalAveragePooling2D, GlobalMaxPooling2D, Reshape, Dense

from tensorflow.keras.layers import multiply, Permute, Concatenate, Conv2D, Add, Activation, Lambda

from tensorflow.keras.layers import Dropout, MultiHeadAttention, LayerNormalization, Reshape

from tensorflow.keras import backend as K

#%% Create and apply the attention model

def attention_block(in_layer, attention_model, ratio=8, residual = False, apply_to_input=True): 

    in_sh = in_layer.shape # dimensions of the input tensor

    in_len = len(in_sh) 

    expanded_axis = 2 # defualt = 2

    if attention_model == 'mha':   # Multi-head self attention layer 

        if(in_len > 3):

            in_layer = Reshape((in_sh[1],-1))(in_layer)

        out_layer = mha_block(in_layer)

    elif attention_model == 'mhla':  # Multi-head local self-attention layer 

        if(in_len > 3):

            in_layer = Reshape((in_sh[1],-1))(in_layer)

        out_layer = mha_block(in_layer, vanilla = False)

    elif attention_model == 'se':   # Squeeze-and-excitation layer

        if(in_len < 4):

            in_layer = tf.expand_dims(in_layer, axis=expanded_axis)

        out_layer = se_block(in_layer, ratio, residual, apply_to_input)

    elif attention_model == 'cbam': # Convolutional block attention module

        if(in_len < 4):

            in_layer = tf.expand_dims(in_layer, axis=expanded_axis)

        out_layer = cbam_block(in_layer, ratio=ratio, residual = residual)

    else:

        raise Exception("'{}' is not supported attention module!".format(attention_model))

    if (in_len == 3 and len(out_layer.shape) == 4):

        out_layer = tf.squeeze(out_layer, expanded_axis)

    elif (in_len == 4 and len(out_layer.shape) == 3):

        out_layer = Reshape((in_sh[1], in_sh[2], in_sh[3]))(out_layer)

    return out_layer

#%% Multi-head self Attention (MHA) block

def mha_block(input_feature, key_dim=8, num_heads=2, dropout = 0.5, vanilla = True):

    """Multi Head self Attention (MHA) block.     

    Here we include two types of MHA blocks: 

            The original multi-head self-attention as described in https://arxiv.org/abs/1706.03762

            The multi-head local self attention as described in https://arxiv.org/abs/2112.13492v1

    """    

    # Layer normalization

    x = LayerNormalization(epsilon=1e-6)(input_feature)

    if vanilla:

        # Create a multi-head attention layer as described in 

        # 'Attention Is All You Need' https://arxiv.org/abs/1706.03762

        x = MultiHeadAttention(key_dim = key_dim, num_heads = num_heads, dropout = dropout)(x, x)

    else:

        # Create a multi-head local self-attention layer as described in 

        # 'Vision Transformer for Small-Size Datasets' https://arxiv.org/abs/2112.13492v1

        # Build the diagonal attention mask

        NUM_PATCHES = input_feature.shape[1]

        diag_attn_mask = 1 - tf.eye(NUM_PATCHES)

        diag_attn_mask = tf.cast([diag_attn_mask], dtype=tf.int8)

        # Create a multi-head local self attention layer.

        # x = MultiHeadAttention_LSA(key_dim = key_dim, num_heads = num_heads, dropout = dropout)(

        #     x, x, attention_mask = diag_attn_mask)

        x = MultiHeadAttention_LSA(key_dim = key_dim, num_heads = num_heads, dropout = dropout)(

            x, x, attention_mask = diag_attn_mask)

    x = Dropout(0.3)(x)

    # Skip connection

    mha_feature = Add()([input_feature, x])

    return mha_feature

#%% Multi head self Attention (MHA) block: Locality Self Attention (LSA)

class MultiHeadAttention_LSA(tf.keras.layers.MultiHeadAttention):

    """local multi-head self attention block

     Locality Self Attention as described in https://arxiv.org/abs/2112.13492v1

     This implementation is taken from  https://keras.io/examples/vision/vit_small_ds/ 

    """    

    def __init__(self, **kwargs):

        super().__init__(**kwargs)

        # The trainable temperature term. The initial value is the square 

        # root of the key dimension.

        self.tau = tf.Variable(math.sqrt(float(self._key_dim)), trainable=True)

    def _compute_attention(self, query, key, value, attention_mask=None, training=None):

        query = tf.multiply(query, 1.0 / self.tau)

        attention_scores = tf.einsum(self._dot_product_equation, key, query)

        attention_scores = self._masked_softmax(attention_scores, attention_mask)

        attention_scores_dropout = self._dropout_layer(

            attention_scores, training=training

        )

        attention_output = tf.einsum(

            self._combine_equation, attention_scores_dropout, value

        )

        return attention_output, attention_scores

#%% Squeeze-and-excitation block

def se_block(input_feature, ratio=8, residual = False, apply_to_input=True):

    """Squeeze-and-Excitation(SE) block.

    As described in https://arxiv.org/abs/1709.01507

    The implementation is taken from https://github.com/kobiso/CBAM-keras

    """

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

    channel = input_feature.shape[channel_axis]

    se_feature = GlobalAveragePooling2D()(input_feature)

    se_feature = Reshape((1, 1, channel))(se_feature)

    assert se_feature.shape[1:] == (1,1,channel)

    if (ratio != 0):

        se_feature = Dense(channel // ratio,

                           activation='relu',

                           kernel_initializer='he_normal',

                           use_bias=True,

                           bias_initializer='zeros')(se_feature)

        assert se_feature.shape[1:] == (1,1,channel//ratio)

    se_feature = Dense(channel,

                       activation='sigmoid',

                       kernel_initializer='he_normal',

                       use_bias=True,

                       bias_initializer='zeros')(se_feature)

    assert se_feature.shape[1:] == (1,1,channel)

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

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

    if(apply_to_input):

        se_feature = multiply([input_feature, se_feature])

    # Residual Connection

    if(residual): 

        se_feature = Add()([se_feature, input_feature])

    return se_feature

#%% Convolutional block attention module

def cbam_block(input_feature, ratio=8, residual = False):

    """ Convolutional Block Attention Module(CBAM) block.

    As described in https://arxiv.org/abs/1807.06521

    The implementation is taken from https://github.com/kobiso/CBAM-keras

    """

    cbam_feature = channel_attention(input_feature, ratio)

    cbam_feature = spatial_attention(cbam_feature)

    # Residual Connection

    if(residual): 

        cbam_feature = Add()([input_feature, cbam_feature])

    return cbam_feature

def channel_attention(input_feature, ratio=8):

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

#     channel = input_feature._keras_shape[channel_axis]

    channel = input_feature.shape[channel_axis]

    shared_layer_one = Dense(channel//ratio,

                             activation='relu',

                             kernel_initializer='he_normal',

                             use_bias=True,

                             bias_initializer='zeros')

    shared_layer_two = Dense(channel,

                             kernel_initializer='he_normal',

                             use_bias=True,

                             bias_initializer='zeros')

    avg_pool = GlobalAveragePooling2D()(input_feature)    

    avg_pool = Reshape((1,1,channel))(avg_pool)

    assert avg_pool.shape[1:] == (1,1,channel)

    avg_pool = shared_layer_one(avg_pool)

    assert avg_pool.shape[1:] == (1,1,channel//ratio)

    avg_pool = shared_layer_two(avg_pool)

    assert avg_pool.shape[1:] == (1,1,channel)

    max_pool = GlobalMaxPooling2D()(input_feature)

    max_pool = Reshape((1,1,channel))(max_pool)

    assert max_pool.shape[1:] == (1,1,channel)

    max_pool = shared_layer_one(max_pool)

    assert max_pool.shape[1:] == (1,1,channel//ratio)

    max_pool = shared_layer_two(max_pool)

    assert max_pool.shape[1:] == (1,1,channel)

    cbam_feature = Add()([avg_pool,max_pool])

    cbam_feature = Activation('sigmoid')(cbam_feature)

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

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

    return multiply([input_feature, cbam_feature])

def spatial_attention(input_feature):

    kernel_size = 7

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

        channel = input_feature.shape[1]

        cbam_feature = Permute((2,3,1))(input_feature)

    else:

        channel = input_feature.shape[-1]

        cbam_feature = input_feature

    avg_pool = Lambda(lambda x: K.mean(x, axis=3, keepdims=True))(cbam_feature)

    assert avg_pool.shape[-1] == 1

    max_pool = Lambda(lambda x: K.max(x, axis=3, keepdims=True))(cbam_feature)

    assert max_pool.shape[-1] == 1

    concat = Concatenate(axis=3)([avg_pool, max_pool])

    assert concat.shape[-1] == 2

    cbam_feature = Conv2D(filters = 1,

                    kernel_size=kernel_size,

                    strides=1,

                    padding='same',

                    activation='sigmoid',

                    kernel_initializer='he_normal',

                    use_bias=False)(concat)    

    assert cbam_feature.shape[-1] == 1

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

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

    return multiply([input_feature, cbam_feature])