import numpy as np
from tensorflow import keras as ks

BATCH_SIZE = 32
NUM_DENSE_LAYERS = 1
NUM_FEATURES = 256
NUM_LSTM_LAYERS = 2
VOCAB_SIZE = 262_144


def l2norm(matrix):
    return ks.backend.sqrt(
        ks.backend.sum(
            ks.backend.square(matrix),
            axis=-1,
            keepdims=True,
        ),
    )


class CosineSimilarity(ks.layers.Layer):
    def __init__(self, activation=None, **kwargs):
        self.activation = ks.activations.get(activation)
        super().__init__(trainable=False, **kwargs)

    def call(self, inputs):
        left, right = (
            matrix / (l2norm(matrix) + ks.backend.epsilon())
            for matrix in inputs
        )
        output = ks.backend.dot(left, ks.backend.transpose(right))
        return self.activation(output)

    def compute_output_shape(self, input_shape):
        shape_left, shape_right = input_shape
        return (shape_left[0], shape_right[0])


class TextPipeline(ks.models.Sequential):
    def __init__(self):
        super().__init__()
        self.add(ks.layers.Embedding(
            input_dim=VOCAB_SIZE,
            output_dim=NUM_FEATURES,
            mask_zero=True,
            name='word_embeddings',
        ))
        for _ in range(NUM_LSTM_LAYERS - 1):
            self.add(ks.layers.LSTM(
                units=NUM_FEATURES,
                return_sequences=True,
            ))
            # self.add(ks.layers.BatchNormalization())
        self.add(ks.layers.LSTM(
            units=NUM_FEATURES,
        ))
        for _ in range(NUM_DENSE_LAYERS):
            # self.add(ks.layers.BatchNormalization())
            self.add(ks.layers.Dense(
                units=NUM_FEATURES,
                activation='tanh',
            ))


class Brain(ks.models.Model):
    def __init__(self):
        offer_input = ks.Input(
            shape=(None,),
            dtype='int32',
            name='offer_input',
        )
        item_input = ks.Input(
            shape=(None,),
            dtype='int32',
            name='item_input',
        )
        text_pipeline = TextPipeline()
        offer_text_embedding = text_pipeline(offer_input)
        item_text_embedding = text_pipeline(item_input)
        similarity = CosineSimilarity(
            # activation='softmax',
            activation='relu',
        )([offer_text_embedding, item_text_embedding])

        super().__init__(
            inputs=[offer_input, item_input],
            outputs=[similarity],
            name='brain',
        )
        self.compile(
            optimizer='adam',
            loss='mean_squared_error',
        )


offers = np.array([[     0,      0,      0,      0],
       [    72,   1536,    408, 120614],
       [   842,  74200,    665,   4342],
       [ 17175,  24987,  13786,   1062],
       [   421,   1439,      6,     63],
       [    72,  23805,  33627,  24301],
       [  2977,      2, 133735,  16077],
       [    72,   2081,  77235,  77234],
       [   390,    347,      2,      2],
       [    11,      2,   6339,      2],
       [ 35016,  98448,     10,  18498],
       [  3654, 149776,   2395,    701],
       [  1413,    361,      4,   6906],
       [   712,   9577,      2,  52794],
       [    11,      2,  34073,      2],
       [ 50408,  12263,   9709,  39344],
       [    11,      2,   2225,      2],
       [    72, 133515,  18916,   7087],
       [     0,      0,      0,      0],
       [   781,     51,      2,    947],
       [    72,   7585,  86920,      2],
       [     0,      0,      0,      0],
       [ 26244,  11879,  17887,  92594],
       [ 35016,  98448,     10,  18498],
       [  4236,  15909,     39,   6955],
       [  6955,     59,     71,   5225],
       [ 50408,   1436,   8509,   6922],
       [ 84705,    141,    106,      2],
       [   390,    347,     14,    209],
       [   426,     78,   1291,    385],
       [   918,      2,    892,      2],
       [  1766,      2,   1539, 146045]], dtype='int32')
items = np.array([[  2777,  10985,   8358,   2777],
       [120614,  10618,  38127,  14625],
       [ 73393,  37841,    535,    842],
       [ 17175,  24987,  13786,   1062],
       [ 14415, 107785, 107783,  89041],
       [ 23805,  33627,  24301,  23805],
       [  2977,      2,   2040,  72082],
       [ 77235,  77234,    793,   7949],
       [100147, 100228,     28,     96],
       [  6339,  65544,  35803,  13368],
       [ 56546,  35016,  98448,  31056],
       [ 34687,   3747,      2,  73159],
       [ 15970,  43435,      2,  43438],
       [ 38108,  77294,  92322,     96],
       [ 34073,   2270,   2093,    121],
       [ 39344,  39356,  12263,   9709],
       [  2225, 122243,     19,   1916],
       [  5076,  60819,  17653,   4958],
       [  6339,  36379,  36470,  67052],
       [  2752,   9196,  20315,      2],
       [  7585,  86920,      2, 113815],
       [  2041,   3800,     79,    728],
       [ 26244,  11879,  17887,  92594],
       [ 56546,  35016,   5651,  98446],
       [ 94436,  20753,      2,   2026],
       [ 24701,  24988,      2,  24998],
       [ 41505,    835,   7559,      2],
       [ 70324,  70323,      2,  70328],
       [  2933,    840,  88005,      2],
       [ 86927,  92144,      2,  92143],
       [  2752,  14990,      2,  18987],
       [  4999,   3191,  85175,    758]], dtype='int32')
diag_matrix = np.diag(np.ones((BATCH_SIZE,)))


def generate_batches():
    while True:
        yield (
            {
                'offer_input': offers,
                'item_input': items,
            },
            diag_matrix,
        )


brain = Brain()
brain.fit_generator(generate_batches(), epochs=1, steps_per_epoch=4)
print(brain.layers[2].layers[3].weights[0].eval(session=ks.backend.get_session()))
# [[nan nan nan ... nan nan nan]
#  [nan nan nan ... nan nan nan]
#  [nan nan nan ... nan nan nan]
#  ...
#  [nan nan nan ... nan nan nan]
#  [nan nan nan ... nan nan nan]
#  [nan nan nan ... nan nan nan]]