import tensorflow as tf
from .encoder import Encoder
from .decoder import Decoder
from tensorflow.keras.layers import Dense

@tf.keras.utils.register_keras_serializable()
class Transformer(tf.keras.Model):
    def __init__(self, num_layers, d_model, num_heads, dff, input_vocab_size, 
                 target_vocab_size, max_tokens, dropout_rate=0.1, **kwargs):
        super(Transformer, self).__init__(**kwargs)
        self.num_layers = num_layers
        self.d_model = d_model
        self.num_heads = num_heads
        self.dff = dff
        self.input_vocab_size = input_vocab_size
        self.target_vocab_size = target_vocab_size
        self.max_tokens = max_tokens
        self.dropout_rate = dropout_rate
        
        self.encoder = Encoder(num_layers, d_model, num_heads, dff, 
                             input_vocab_size, max_tokens, dropout_rate)
        self.decoder = Decoder(num_layers, d_model, num_heads, dff, 
                              target_vocab_size, max_tokens, dropout_rate)
        self.final_layer = Dense(target_vocab_size)

    def call(self, inputs, training=None):
        enc_input, dec_input = inputs
        enc_padding_mask = self.create_padding_mask(enc_input)
        look_ahead_mask = self.create_look_ahead_mask(tf.shape(dec_input)[1])
        dec_padding_mask = self.create_padding_mask(enc_input)
        enc_output = self.encoder(enc_input, training=training, mask=enc_padding_mask)
        dec_output = self.decoder(dec_input, enc_output, training=training, 
                                 look_ahead_mask=look_ahead_mask, 
                                 padding_mask=dec_padding_mask)
        final_output = self.final_layer(dec_output)
        return final_output

    def create_padding_mask(self, seq):
        mask = tf.cast(tf.math.equal(seq, 0), tf.float32)
        return mask[:, tf.newaxis, tf.newaxis, :]

    def create_look_ahead_mask(self, size):
        mask = 1 - tf.linalg.band_part(tf.ones((size, size)), -1, 0)
        return mask

    def get_config(self):
        config = super().get_config()
        config.update({
            'num_layers': self.num_layers,
            'd_model': self.d_model,
            'num_heads': self.num_heads,
            'dff': self.dff,
            'input_vocab_size': self.input_vocab_size,
            'target_vocab_size': self.target_vocab_size,
            'max_tokens': self.max_tokens,
            'dropout_rate': self.dropout_rate
        })
        return config