class Transformer(nn.Module):
    '''
    define the whole architecture of Transformer in:
        Vaswani et al. Attention is All You Need. NIPS 2017.
    '''
    def __init__(self, num_heads=8, d_model=512, d_ff=2048, num_enc_layers=6, num_dec_layers=6,
                 drop_rate=0.1, warmup_steps=400, pos_encoding_type='sinusoid',
                 ls_epsilon=0.1, use_label_smoothing=True,
                 model_name='transformer', tf_sess_config=None, **kwargs):
        super().__init__()
        self.h = num_heads
        self.d_model = d_model
        self.d_ff = d_ff

        self.num_enc_layers = num_enc_layers
        self.num_dec_layers = num_dec_layers

        # Dropout regularization: added in every sublayer before layer_norm(...) and
        # applied to embedding + positional encoding.
        self.drop_rate = drop_rate

        # Label smoothing epsilon
        self.ls_epsilon = ls_epsilon
        self.use_label_smoothing = use_label_smoothing
        self.pos_encoding_type = pos_encoding_type

        # For computing the learning rate
        self.warmup_steps = warmup_steps

        self.config = dict(
            num_heads=self.h,
            d_model=self.d_model,
            d_ff=self.d_ff,
            num_enc_layers=self.num_enc_layers,
            num_dec_layers=self.num_dec_layers,
            drop_rate=self.drop_rate,
            warmup_steps=self.warmup_steps,
            ls_epsilon=self.ls_epsilon,
            use_label_smoothing=self.use_label_smoothing,
            pos_encoding_type=self.pos_encoding_type,
            model_name=self.model_name,
            tf_sess_config=self.tf_sess_config,
        )
        def forward(self, src, tgt, src_mask, tgt_mask): # 这里进行拼接，将encoder和decoder两大模块拼接在一起
            enc_out = self.encoder(src, src_mask)
            dec_out = self.decoder(enc_out, src_mask, tgt, tgt_mask)
            return dec_out
        

def clones(module, N):
    "Produce N identical layers."
    return nn.ModuleList([copy.deepcopy(module) for _ in range(N)])

class TransformerEncoder(nn.Module):
    def __init__(self, encoder_layer, num_enc_layers) -> None:
        self.num_enc_layers = num_enc_layers
        self.encoder_layers = clones(encoder_layer,num_enc_layers) # 将encoder_layer复制6次

    def forward(self, src, src_mask):
        out = src
        for layer in self.encoder_layers:
            out = layer(out, src_mask)
        return out

out = inp  # now, (batch, seq_len, embed_size)
with tf.variable_scope(scope):
    for i in range(self.num_enc_layers):
        out = self.encoder_layer(out, input_mask, f'enc_{i}')
        return out

class TransformerEncoderLayer(nn.Module):
    """
    Args:
        d_model: the number of expected features in the input (required).
        nhead: the number of heads in the multiheadattention models (required).
        dim_feedforward: the dimension of the feedforward network model (default=2048).
    About:

    """
    # One multi-head attention + one feed-forward
    def __init__(self, d_model, n_head, dim_feedforward, dropout = 0.1) -> None:
        super().__init__()
        self.self_attn = MultiheadAttention(d_model, n_head)
        self.norm_1 = nn.LayerNorm(d_model)
        # Implementation of Feedforward model(Two linear transformation together with one dropout)
        self.linear1 = nn.Linear(d_model, dim_feedforward, )
        self.dropout = nn.Dropout(dropout)
        self.linear2 = nn.Linear(dim_feedforward, d_model)
        self.norm_2 = nn.LayerNorm(d_model)

    def __ff_block(self, x):
        # feed forward layer contains two linear
        out = F.relu(self.linear1(x))
        out = self.dropout(out)
        out = self.linear2(out)
        
    def forward(self, src, src_mask):
        out = src
        out = self.norm_1(out + self.self_attn(out, src_mask))# 这里在pytorch的官方实现中在self_attention后还加了一个dropout
        out = self.norm_2(out + self.__ff_block(out))
    
        return out

def multihead_attention(self, query, memory=None, mask=None, scope='attn'):
        """
        Args:
            query (tf.tensor): of shape (batch, q_size, d_model)
            memory (tf.tensor): of shape (batch, m_size, d_model)
            mask (tf.tensor): shape (batch, q_size, k_size)

        Returns:h
            a tensor of shape (bs, q_size, d_model)
        """
        if memory is None:
            memory = query

        with tf.variable_scope(scope):
            # Linear project to d_model dimension: [batch, q_size/k_size, d_model]
            Q = tf.layers.dense(query, self.d_model, activation=tf.nn.relu)
            K = tf.layers.dense(memory, self.d_model, activation=tf.nn.relu)
            V = tf.layers.dense(memory, self.d_model, activation=tf.nn.relu)

            # Split the matrix to multiple heads and then concatenate to have a larger
            # batch size: [h*batch, q_size/k_size, d_model/num_heads]
            Q_split = tf.concat(tf.split(Q, self.h, axis=2), axis=0)
            K_split = tf.concat(tf.split(K, self.h, axis=2), axis=0)
            V_split = tf.concat(tf.split(V, self.h, axis=2), axis=0)
            mask_split = tf.tile(mask, [self.h, 1, 1])

            # Apply scaled dot product attention
            out = self.scaled_dot_product_attention(Q_split, K_split, V_split, mask=mask_split)

            # Merge the multi-head back to the original shape
            out = tf.concat(tf.split(out, self.h, axis=0), axis=2)  # [bs, q_size, d_model]

            # The final linear layer and dropout.
            # out = tf.layers.dense(out, self.d_model)
            # out = tf.layers.dropout(out, rate=self.drop_rate, training=self._is_training)

        return out