import tensorflow as tf

try:
    import tensorflow.python.keras as keras
    from tensorflow.python.keras import layers
    from tensorflow.python.keras import backend as K
except:
    import tensorflow.keras as keras
    from tensorflow.keras import layers
    from tensorflow.keras import backend as K


class Attention_layer(layers.Layer):
    def __init__(self, **kwargs):
        super(Attention_layer, self).__init__(**kwargs)

    def build(self, inputs_shape):
        assert isinstance(inputs_shape, list)

        # inputs: --> size
        #     [(None,max_contents,code_vector_size),(None,max_contents)]
        #     the second input is optional
        if (len(inputs_shape) < 1 or len(inputs_shape) > 2):
            raise ValueError("AttentionLayer expect one or two inputs.")

        # (None,max_contents,code_vector_size)
        input_shape = inputs_shape[0]

        if (len(input_shape) != 3):
            raise ValueError("Input shape for AttentionLayer shoud be of 3 dimensions.")

        self.input_length = int(input_shape[1])
        self.input_dim = int(input_shape[2])

        attention_param_shape = (self.input_dim, 1)

        self.attention_param = self.add_weight(
            name='attention_param',
            shape=attention_param_shape,
            initializer='uniform',
            trainable=True,
            dtype=tf.float32
        )

        super(Attention_layer, self).build(input_shape)

    def call(self, inputs, **kwargs):
        assert isinstance(inputs, list)

        # inputs:
        #     [(None,max_contents,code_vector_size),(None,max_contents)]
        #     the second input is optional
        if (len(inputs) < 1 or len(inputs) > 2):
            raise ValueError("AttentionLayer expect one or two inputs.")

        actual_input = inputs[0]
        mask = inputs[1] if (len(inputs) > 1) else None

        if mask is not None and not (((len(mask.shape) == 3 and mask.shape[2] == 1) or (len(mask.shape) == 2)) and (
                mask.shape[1] == self.input_length)):
            raise ValueError(
                "`mask` shoud be of shape (batch, input_length) or (batch, input_length, 1) when calling AttentionLayer.")

        assert actual_input.shape[-1] == self.attention_param.shape[0]

        # (batch, input_length, input_dim) * (input_dim, 1) ==> (batch, input_length, 1)
        attention_weights = K.dot(actual_input, self.attention_param)

        if mask is not None:
            if (len(mask.shape) == 2):
                mask = K.expand_dims(mask, axis=2)  # (batch, input_dim, 1)
            mask = K.log(mask)  # e.g. K.exp(K.log(0.)) = 0 K.exp(K.log(1.)) =1
            attention_weights += mask

        attention_weights = K.softmax(attention_weights, axis=1)
        result = K.sum(actual_input * attention_weights, axis=1)
        return result, attention_weights

    def compute_output_shape(self, input_shape):
        return input_shape[0], input_shape[2]  # (batch,input_length,input_dim) --> (batch,input_dim)