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, 11, 53, 43), 'tensorflow.cast', 'tf.cast', (), '', True, 'import tensorflow as tf\n'), ((55, 11, 55, 32), 'tensorflow.reduce_mean', 'tf.reduce_mean', ({(55, 26, 55, 31): 'loss_'}, {}), '(loss_)', True, 'import tensorflow as tf\n'), ((58, 11, 58, 30), 'numpy.zeros', 'np.zeros', ({(58, 20, 58, 29): 'seq.shape'}, {}), '(seq.shape)', True, 'import numpy as np\n'), ((60, 10, 60, 37), 'tensorflow.math.multiply', 'tf.math.multiply', ({(60, 27, 60, 30): 'seq', (60, 32, 60, 36): 'mask'}, {}), '(seq, mask)', True, 'import tensorflow as tf\n'), ((62, 10, 62, 37), 'tensorflow.math.multiply', 'tf.math.multiply', ({(62, 27, 62, 30): 'seq', (62, 32, 62, 36): 'mask'}, {}), '(seq, mask)', True, 'import tensorflow as tf\n'), ((31, 26, 33, 87), 'tensorflow.keras.layers.Embedding', 'tf.keras.layers.Embedding', (), '', True, 'import tensorflow as tf\n'), ((51, 31, 51, 53), 'tensorflow.math.equal', 'tf.math.equal', ({(51, 45, 51, 49): 'real', (51, 51, 51, 52): '0'}, {}), '(real, 0)', True, 'import tensorflow as tf\n'), ((15, 23, 15, 86), 'tensorflow.keras.layers.Dense', 'tf.keras.layers.Dense', (), '', True, 'import tensorflow as tf\n'), ((35, 23, 35, 44), 'tensorflow.keras.layers.Add', 'tf.keras.layers.Add', ({}, {}), '()', True, 'import tensorflow as tf\n'), ((37, 23, 37, 99), 'tensorflow.keras.layers.Dense', 'tf.keras.layers.Dense', (), '', True, 'import tensorflow as tf\n'), ((10, 53, 14, 86), 'tensorflow.keras.layers.GRU', 'tf.keras.layers.GRU', (), '', True, 'import tensorflow as tf\n')]