# -*- coding: utf-8 -*-
"""
encoder for the transformer:
6 layers.each layers has two sub-layers.
the first is multi-head self-attention mechanism;
the second is position-wise fully connected feed-forward network.
for each sublayer. use LayerNorm(x+Sublayer(x)). all dimension=512.
"""
import tensorflow as tf
from model.base_model import BaseClass
import time
class Encoder(BaseClass):
    def __init__(self,d_model,d_k,d_v,sequence_length,h,batch_size,num_layer,Q,K_s,mask=None,dropout_keep_prob=0.9,use_residual_conn=True):
        """
        :param d_model:
        :param d_k:
        :param d_v:
        :param sequence_length:
        :param h:
        :param batch_size:
        :param embedded_words: shape:[batch_size*sequence_length,embed_size]
        """
        super(Encoder, self).__init__(d_model,d_k,d_v,sequence_length,h,batch_size,num_layer=num_layer)
        self.Q=Q
        self.K_s=K_s
        self.mask=mask
        self.initializer = tf.random_normal_initializer(stddev=0.1)
        self.dropout_keep_prob=dropout_keep_prob
        self.use_residual_conn=use_residual_conn

    def encoder_fn(self):
        """
        use transformer encoder to encode the input, output a sequence. input: [batch_size,sequence_length,d_embedding]
        :return:  output:[batch_size*sequence_length,d_model]
        """
        start = time.time()
        #print("encoder_fn.started.")
        x=self.Q
        for layer_index in range(self.num_layer):
            x=self.encoder_single_layer(x,x,x,layer_index) # Q,K_s,V_s
            #print("encoder_fn.",layer_index,".x:",x)
        end = time.time()
        #print("encoder_fn.ended.x:",x)
        #print("time spent:",(end-start))
        return x

    def encoder_single_layer(self,Q,K_s,V_s,layer_index):
        """
        singel layer for encoder.each layers has two sub-layers:
        the first is multi-head self-attention mechanism; the second is position-wise fully connected feed-forward network.
        for each sublayer. use LayerNorm(x+Sublayer(x)). input and output of last dimension: d_model
        :param Q: shape should be:       [batch_size,sequence_length,d_model]
        :param K_s: shape should be:     [batch_size,sequence_length,d_model]
        :return:output: shape should be: [batch_size,sequence_length,d_model]
        """
        #1.1 the first is multi-head self-attention mechanism
        multi_head_attention_output=self.sub_layer_multi_head_attention(layer_index,Q,K_s,V_s,mask=self.mask,dropout_keep_prob=self.dropout_keep_prob) #[batch_size,sequence_length,d_model]
        #1.2 use LayerNorm(x+Sublayer(x)). all dimension=512.
        multi_head_attention_output=self.sub_layer_layer_norm_residual_connection(K_s,multi_head_attention_output,layer_index,
                                                                                  dropout_keep_prob=self.dropout_keep_prob,use_residual_conn=self.use_residual_conn,sub_layer_name='layer1')
        #2.1 the second is position-wise fully connected feed-forward network.
        postion_wise_feed_forward_output=self.sub_layer_postion_wise_feed_forward(multi_head_attention_output,layer_index)
        #2.2 use LayerNorm(x+Sublayer(x)). all dimension=512.
        postion_wise_feed_forward_output= self.sub_layer_layer_norm_residual_connection(multi_head_attention_output,postion_wise_feed_forward_output,layer_index,
                                                                                        dropout_keep_prob=self.dropout_keep_prob,sub_layer_name='layer2')
        return  postion_wise_feed_forward_output #,postion_wise_feed_forward_output


def init():
    #1. assign value to fields
    vocab_size=1000
    d_model = 512
    d_k = 64
    d_v = 64
    sequence_length = 5*10
    h = 8
    batch_size=4*32
    initializer = tf.random_normal_initializer(stddev=0.1)
    # 2.set values for Q,K,V
    vocab_size=1000
    embed_size=d_model
    Embedding = tf.get_variable("Embedding_E", shape=[vocab_size, embed_size],initializer=initializer)
    input_x = tf.placeholder(tf.int32, [batch_size,sequence_length], name="input_x") #[4,10]
    print("input_x:",input_x)
    embedded_words = tf.nn.embedding_lookup(Embedding, input_x) #[batch_size*sequence_length,embed_size]
    Q = embedded_words  # [batch_size*sequence_length,embed_size]
    K_s = embedded_words  # [batch_size*sequence_length,embed_size]
    V_s = embedded_words  # [batch_size*sequence_length,embed_size]
    num_layer=6
    mask = get_mask(batch_size, sequence_length)
    #3. get class object
    encoder_class=Encoder(d_model,d_k,d_v,sequence_length,h,batch_size,num_layer,Q,K_s,mask=mask) #Q,K_s,embedded_words
    return encoder_class,Q,K_s,V_s

def get_mask(batch_size,sequence_length):
    lower_triangle=tf.matrix_band_part(tf.ones([sequence_length,sequence_length]),-1,0)
    result=-1e9*(1.0-lower_triangle)
    print("get_mask==>result:",result)
    return result

def test_postion_wise_feed_forward(encoder_class,x,layer_index):
    sub_layer_postion_wise_feed_forward_output=encoder_class.sub_layer_postion_wise_feed_forward(x, layer_index)
    return sub_layer_postion_wise_feed_forward_output

def test_sub_layer_multi_head_attention(encoder_class,index_layer,Q,K_s,V_s):
    sub_layer_multi_head_attention_output=encoder_class.sub_layer_multi_head_attention(index_layer,Q,K_s,V_s)
    return sub_layer_multi_head_attention_output


encoder_class,Q,K_s,V_s=init()


#below is 4 callable codes for testing functions: from sub(small) function to whole function of encoder.

def test():
    #1.test 1: for sub layer of multi head attention
    index_layer=0
    #sub_layer_multi_head_attention_output=test_sub_layer_multi_head_attention(encoder_class,index_layer,Q,K_s,V_s)
    #print("sub_layer_multi_head_attention_output1:",sub_layer_multi_head_attention_output)

    #2. test 2: for sub layer of multi head attention with poistion-wise feed forward
    #d1,d2,d3=sub_layer_multi_head_attention_output.get_shape().as_list()
    #print("d1:",d1,";d2:",d2,";d3:",d3)
    #postion_wise_ff_input=sub_layer_multi_head_attention_output #tf.reshape(sub_layer_multi_head_attention_output,shape=[-1,d3])
    #print("sub_layer_postion_wise_feed_forward_input:",postion_wise_ff_input)
    #sub_layer_postion_wise_feed_forward_output=test_postion_wise_feed_forward(encoder_class,postion_wise_ff_input,index_layer)
    #sub_layer_postion_wise_feed_forward_output=tf.reshape(sub_layer_postion_wise_feed_forward_output,shape=(d1,d2,d3))
    #print("sub_layer_postion_wise_feed_forward_output:",sub_layer_postion_wise_feed_forward_output)
    #3.test 3: test for single layer of encoder
    #encoder_class.encoder_single_layer(Q,K_s,V_s,index_layer)
    #4.test 4: test for encoder. with N layers

    representation = encoder_class.encoder_fn()
    print("representation:",representation)

# test()