import numpy as np
import tensorflow as tf

class Bidirectional(tf.keras.Model):
    def __init__(self, units: int, 
                 projection_units: int):
        super(Bidirectional, self).__init__()
        self.units = units
        self.projection_units = projection_units
        self.Layers = [tf.keras.layers.Bidirectional(tf.keras.layers.GRU(self.units, 
                                                                         return_sequences=True, 
                                                                         return_state=True, 
                                                                         recurrent_initializer='glorot_uniform', 
                                                                         name='birnn')),
                       tf.keras.layers.Dense(self.projection_units, name='projection')]

    def call(self, inp):
        out, _, _ = self.Layers[0](inp)
        out = self.Layers[1](out)
        return out

class BiRNN(tf.keras.Model):
    def __init__(self, units: int,projection_units: int,max_seq_length: int,
                 vocab_size: int,embedding_dim: int,embedding_matrix = None):
        super(BiRNN, self).__init__()
        self.units = units
        self.projection_units=projection_units
        self.max_seq_length = max_seq_length
        self.vocab_size = vocab_size
        self.embedding_dim = embedding_dim
        self.embeddings = tf.keras.layers.Embedding(self.vocab_size, self.embedding_dim,
                                                    weights = [embedding_matrix], 
                                                    trainable=False, name='embeddings')
        self.Layers = [Bidirectional(units=self.units, projection_units=self.projection_units),
                       tf.keras.layers.Add(),
                       Bidirectional(units=self.units, projection_units=self.projection_units),
                       tf.keras.layers.Dense(self.vocab_size, activation='softmax', name='softmax')]

    def call(self, inp, predict=False):
        inp = self.embeddings(inp)
        out1 = self.Layers[0](inp)
        out2 = self.Layers[1]([inp, out1])
        out3 = self.Layers[2](out2)
        if predict is False:
            return out3
        else:
            out4 = self.Layers[3](out3)
            return out4

def loss_function(real, pred, loss_object):
    mask = tf.math.logical_not(tf.math.equal(real, 0))
    loss_ = loss_object(real, pred)
    mask = tf.cast(mask, dtype=loss_.dtype)
    loss_ *= mask
    return tf.reduce_mean(loss_)

def mask_sequences(seq, t):
    mask = np.zeros(seq.shape)
    mask[:,:t] = 1
    inp = tf.math.multiply(seq, mask)
    mask[:,:t+1] = 1
    tar = tf.math.multiply(seq, mask)

    return inp, tar
[(53, 'tensorflow.cast', 'tf.cast', True, 'import tensorflow as tf\n'), (55, 'tensorflow.reduce_mean', 'tf.reduce_mean', True, 'import tensorflow as tf\n'), (58, 'numpy.zeros', 'np.zeros', True, 'import numpy as np\n'), (60, 'tensorflow.math.multiply', 'tf.math.multiply', True, 'import tensorflow as tf\n'), (62, 'tensorflow.math.multiply', 'tf.math.multiply', True, 'import tensorflow as tf\n'), (31, 'tensorflow.keras.layers.Embedding', 'tf.keras.layers.Embedding', True, 'import tensorflow as tf\n'), (51, 'tensorflow.math.equal', 'tf.math.equal', True, 'import tensorflow as tf\n'), (15, 'tensorflow.keras.layers.Dense', 'tf.keras.layers.Dense', True, 'import tensorflow as tf\n'), (35, 'tensorflow.keras.layers.Add', 'tf.keras.layers.Add', True, 'import tensorflow as tf\n'), (37, 'tensorflow.keras.layers.Dense', 'tf.keras.layers.Dense', True, 'import tensorflow as tf\n'), (10, 'tensorflow.keras.layers.GRU', 'tf.keras.layers.GRU', True, 'import tensorflow as tf\n')]