import tensorflow as tf
from tensorflow import keras

@keras.utils.register_keras_serializable()
class BlockDecoders (keras.Layer) :
  def __init__ (self,d_model,ffn_dim,num_heads,dropout_rate=0.1,**kwargs) :
    super(BlockDecoders,self).__init__(**kwargs)
    self.mha = keras.layers.MultiHeadAttention(num_heads=num_heads,key_dim=d_model,dropout=dropout_rate)
    self.normal1 = keras.layers.LayerNormalization(epsilon=1e-6)
    self.ffn = keras.Sequential([
        keras.layers.Dense(ffn_dim,activation=keras.activations.gelu),
        keras.layers.Dense(d_model)
    ])
    self.dropout = keras.layers.Dropout(rate=dropout_rate)
    self.normal2 = keras.layers.LayerNormalization(epsilon=1e-6)
    self.d_model = d_model
    self.ffn_dim = ffn_dim
    self.num_head = num_heads
    self.dropout_rate = dropout_rate

  def call(self,x,training=False) :
    attn = self.mha(x,x,x,training=training,use_causal_mask=True)
    attn = self.normal1(attn + x)

    ffn = self.ffn(attn)
    ffn = self.dropout(ffn)
    ffn = self.normal2(ffn + attn)
    return ffn

  def get_config(self) :
    config = super(BlockDecoders,self).get_config()
    config.update ({
        'd_model' : self.d_model,
        'ffn_dim' : self.ffn_dim,
        'num_head' : self.num_head,
        'dropout_rate' : self.dropout_rate
    })
    return config

  @classmethod
  def from_config(cls,config) :
    return cls(**config)

@keras.utils.register_keras_serializable()
class Micro_Gen_Teks (keras.Model) :
  def __init__ (self,vocab_size,d_model,ffn_dim,num_heads,num_blocks,maxpos,dropout_rate=0.1,**kwargs) :
    super(Micro_Gen_Teks,self).__init__(**kwargs)
    self.Embedding = keras.layers.Embedding(vocab_size,d_model)
    self.pos_embedding = keras.layers.Embedding(maxpos,d_model)
    self.BlockDecoders = [BlockDecoders(
        d_model=d_model,ffn_dim=ffn_dim,num_heads=num_heads,dropout_rate=dropout_rate
    ) for _ in range(num_blocks)]
    self.final_layer = keras.layers.Dense(vocab_size)

    self.vocab_size = vocab_size
    self.d_model = d_model
    self.ffn_dim = ffn_dim
    self.num_heads = num_heads
    self.num_blocks = num_blocks
    self.dropout_rate = dropout_rate
    self.maxpos = maxpos

  def call(self,x,training = True) :
    batch, seq = tf.shape(x)[0], tf.shape(x)[1]
    pos = tf.range(start=0,limit=seq,delta=1)
    pos = self.pos_embedding(pos)
    pos = tf.expand_dims(pos,axis=0)
    x = self.Embedding(x)
    x *= tf.sqrt(tf.cast(self.d_model,dtype=tf.float16))
    x = x + pos
    for block in self.BlockDecoders :
      x = block(x,training=training)

    x = self.final_layer(x)
    return x


  def get_config(self) :
    config = super(Micro_Gen_Teks,self).get_config()
    config.update({
        'vocab_size' : self.vocab_size,
        'd_model' : self.d_model,
        'ffn_dim' : self.ffn_dim,
        'num_heads' : self.num_heads,
        'num_blocks' : self.num_blocks,
        'dropout_rate' : self.dropout_rate,
        'maxpos' : self.maxpos
    })
    return config

  @classmethod
  def from_config(cls,config) :
    return cls(**config)