from tensorflow.python.keras import backend as K
from tensorflow.keras.models import Model
from tensorflow.keras.layers import Input
from tensorflow.keras.layers import Layer
from tensorflow.keras.layers import LSTM
from tensorflow.keras.layers import Embedding
from tensorflow.python.keras.layers.wrappers import Bidirectional
from tensorflow.keras.models import Sequential
from tensorflow.keras.optimizers import Adam

class ManDist(Layer):
    def __init__(self, **kwargs):
        self.result = None
        super(ManDist, self).__init__(**kwargs)

    def build(self, input_shape):
        super(ManDist, self).build(input_shape)

    def call(self, x, **kwargs):
        self.result = K.exp(-K.sum(K.abs(x[0] - x[1]), axis=1, keepdims=True))
        return self.result

    def compute_output_shape(self):
        return K.int_shape(self.result)

def build_siamese_model(embedding_matrix, embeddingDim, max_sequence_length=384, number_lstm_units=50, rate_drop_lstm=0.01):

    x = Sequential()
    x.add(Embedding(len(embedding_matrix), embeddingDim, weights=[embedding_matrix], input_shape=(max_sequence_length,), trainable=False))
    x.add(LSTM(number_lstm_units, dropout=rate_drop_lstm, return_sequences=True, activation='softmax'))

    input_1 = Input(shape=(max_sequence_length,), dtype='int32')
    input_2 = Input(shape=(max_sequence_length,), dtype='int32')

    distance = ManDist()([x(input_1), x(input_2)])
    model = Model(inputs=[input_1, input_2], outputs=[distance])
    model.compile(loss='mean_squared_error', optimizer=Adam(learning_rate=0.001), metrics=['accuracy'])

    return model