# Copyright (c) 2017-present, Facebook, Inc.
# All rights reserved.
#
# This source code is licensed under the license found in the LICENSE file in
# the root directory of this source tree. An additional grant of patent rights
# can be found in the PATENTS file in the same directory.

import math

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.utils.checkpoint as checkpoint
from torch.nn import Parameter

from fairseq.modules import (
    LearnedPositionalEmbedding, MultiheadAttention,
    SinusoidalPositionalEmbedding,
)

from fairseq.modules.layer_stacker import CreateLayerStack
from fairseq.modules.outer_layer_stacker import CreateOuterLayerStack

from . import (
    FairseqIncrementalDecoder, FairseqEncoder, FairseqModel,
    register_model, register_model_architecture, transformer
)

@register_model('stack_transformer')
class StackTransformerModel(FairseqModel):
    def __init__(self, encoder, decoder):
        super().__init__(encoder, decoder)

    @staticmethod
    def add_args(parser):
        """Add model-specific arguments to the parser."""
        parser.add_argument('--dropout', type=float, metavar='D',
                            help='dropout probability')
        parser.add_argument('--attention-dropout', type=float, metavar='D',
                            help='dropout probability for attention weights')
        parser.add_argument('--relu-dropout', type=float, metavar='D',
                            help='dropout probability after ReLU in FFN')
        parser.add_argument('--encoder-embed-dim', type=int, metavar='N',
                            help='encoder embedding dimension')
        parser.add_argument('--encoder-ffn-embed-dim', type=int, metavar='N',
                            help='encoder embedding dimension for FFN')
        parser.add_argument('--encoder-layers', type=int, metavar='N',
                            help='num encoder layers')
        parser.add_argument('--encoder-attention-heads', type=int, metavar='N',
                            help='num encoder attention heads')
        parser.add_argument('--encoder-normalize-before', default=True, action='store_true',
                            help='apply layernorm before each encoder block')
        parser.add_argument('--encoder-learned-pos', default=False, action='store_true',
                            help='use learned positional embeddings in the encoder')
        parser.add_argument('--decoder-embed-dim', type=int, metavar='N',
                            help='decoder embedding dimension')
        parser.add_argument('--decoder-ffn-embed-dim', type=int, metavar='N',
                            help='decoder embedding dimension for FFN')
        parser.add_argument('--decoder-layers', type=int, metavar='N',
                            help='num decoder layers')
        parser.add_argument('--decoder-attention-heads', type=int, metavar='N',
                            help='num decoder attention heads')
        parser.add_argument('--decoder-learned-pos', default=False, action='store_true',
                            help='use learned positional embeddings in the decoder')
        parser.add_argument('--decoder-normalize-before', default=True, action='store_true',
                            help='apply layernorm before each decoder block')
        parser.add_argument('--share-decoder-input-output-embed', default=False, action='store_true',
                            help='share decoder input and output embeddings')
        parser.add_argument('--share-all-embeddings', default=False, action='store_true',
                            help='share encoder, decoder and output embeddings'
                                 ' (requires shared dictionary and embed dim)')
        ### stack layer parameters
        parser.add_argument('--encoder-inner-stack-type', type=str, choices=['residual', 'dense', 'gru', 'mhsa', 'aan', 'None'],
                            help='the mode of stacking each layer in encoder')
        parser.add_argument('--decoder-inner-stack-type', type=str, choices=['residual', 'dense', 'gru', 'mhsa', 'aan', 'None'],
                            help='the mode of stacking each layer in decoder')
        parser.add_argument('--encoder-outer-stack-type', type=str, choices=['residual', 'dense', 'gru', 'mhsa', 'aan', 'None'],
                            help='the mode of stacking each layer in encoder')
        parser.add_argument('--decoder-outer-stack-type', type=str, choices=['residual', 'dense', 'gru', 'mhsa', 'aan', 'None'],
                            help='the mode of stacking each layer in decoder')
        parser.add_argument('--encoder-layer-norm-owned-type', type=str, choices=['every_sublayer',
                                                                                  'every_block',
                                                                                  'every_same_sublayer',
                                                                                  'every_side'],
                            help='type of owning layer norm in encoder. '
                                 'assume that block_size=N, each block has X sublayers,'
                                 'every_sublayer=N*X, every_block=N, every_same_sublayer=X, every_side=1'
                                 'every_sublayer is used for Standard Transformer')
        parser.add_argument('--decoder-layer-norm-owned-type', type=str, choices=['every_sublayer',
                                                                                  'every_block',
                                                                                  'every_side'],
                            help='type of owning layer norm in decoder. every_sublayer is used for Standard Transformer')
        parser.add_argument('--stack-dropout', type=float, metavar='D',
                            help='push stack dropout probability')
        parser.add_argument('--stack-relu-dropout', type=float, metavar='D',
                            help='stack relu dropout probability')
        parser.add_argument('--stack-att-dropout', type=float, metavar='D',
                            help='stack attention dropout probability')
        parser.add_argument('--encoder-gru-owned-type', type=str, default='every_side',
                            choices=['every_sublayer','every_block','every_side'],
                            help='type of owning gru in encoder GRU Stack')
        parser.add_argument('--decoder-gru-owned-type', type=str, default='every_side',
                            choices=['every_sublayer', 'every_block', 'every_side'],
                            help='type of owning gru in decoder GRU Stack')
        parser.add_argument('--add-layer-embed', default=True, action='store_true',
                            help='apply layer embedding for each stacked layer')
        parser.add_argument('--stack-head-num', type=int, metavar='N',
                            help='head num in MHSA Stack')

    @classmethod
    def build_model(cls, args, task):
        """Build a new model instance."""
        src_dict, tgt_dict = task.source_dictionary, task.target_dictionary

        def build_embedding(dictionary, embed_dim):
            num_embeddings = len(dictionary)
            padding_idx = dictionary.pad()
            return Embedding(num_embeddings, embed_dim, padding_idx)

        if args.share_all_embeddings:
            if src_dict != tgt_dict:
                raise RuntimeError('--share-all-embeddings requires a joined dictionary')
            if args.encoder_embed_dim != args.decoder_embed_dim:
                raise RuntimeError(
                    '--share-all-embeddings requires --encoder-embed-dim to match --decoder-embed-dim')
            encoder_embed_tokens = build_embedding(src_dict, args.encoder_embed_dim)
            decoder_embed_tokens = encoder_embed_tokens
            args.share_decoder_input_output_embed = True
        else:
            encoder_embed_tokens = build_embedding(src_dict, args.encoder_embed_dim)
            decoder_embed_tokens = build_embedding(tgt_dict, args.decoder_embed_dim)

        encoder = StackTransformerEncoder(args, src_dict, encoder_embed_tokens)
        decoder = StackTransformerDecoder(args, tgt_dict, decoder_embed_tokens)
        return StackTransformerModel(encoder, decoder)


class StackTransformerEncoder(FairseqEncoder):
    """Stack Transformer encoder."""

    def __init__(self, args, dictionary, embed_tokens, left_pad=True):
        super().__init__(dictionary)
        self.dropout = args.dropout

        embed_dim = embed_tokens.embedding_dim
        self.padding_idx = embed_tokens.padding_idx

        self.embed_tokens = embed_tokens
        self.embed_scale = math.sqrt(embed_dim)
        self.embed_positions = PositionalEmbedding(
            1024, embed_dim, self.padding_idx,
            left_pad=left_pad,
            learned=args.encoder_learned_pos,
        )
        self.inner_stack_list = nn.ModuleList([])
        for i in range(args.encoder_layers):
            self.inner_stack_list.append(CreateLayerStack(args, True,True))
        self.outer_stack = CreateOuterLayerStack(args, True)
        # self.stack = CreateLayerStack(args, True)
        self.layers = nn.ModuleList([])
        self.layers.extend([
            TransformerEncoderLayer(args)
            for i in range(args.encoder_layers)])
        if args.encoder_outer_stack_type =="transparent":
            self.transparent = True
        else:
            self.transparent = False
    def forward(self, src_tokens, src_lengths):
        for i in range(len(self.inner_stack_list)):
            self.inner_stack_list[i].clear()
        self.outer_stack.clear()
        # embed tokens and positions
        x = self.embed_scale * self.embed_tokens(src_tokens)
        x += self.embed_positions(src_tokens)
        x = F.dropout(x, p=self.dropout, training=self.training)
        # B x T x C -> T x B x C
        x = x.transpose(0, 1).contiguous()
        #print('continuous:{}'.format(x.is_contiguous()))
        self.inner_stack_list[0].push(x,use_dropout=False)
        self.outer_stack.push(x,use_dropout=False)

        # compute padding mask
        encoder_padding_mask = src_tokens.eq(self.padding_idx)
        if not encoder_padding_mask.any():
            encoder_padding_mask = None

        # encoder layers
        for i,layer in enumerate(self.layers):
            self.inner_stack_list[i]=layer(self.inner_stack_list[i], encoder_padding_mask)
            inner_out = self.inner_stack_list[i].pop()
            self.outer_stack.push(inner_out,use_dropout=False)
            self.inner_stack_list[i].clear()
            if i != len(self.layers) - 1:
                self.inner_stack_list[i+1].push(self.outer_stack.pop(), use_dropout=False)

        x = self.outer_stack.pop()
        encoder_list = []
        if self.transparent:
            for i in range(len(self.outer_stack.stack)):
                encoder_list.append(self.outer_stack.stack[i])

            B, L, H = encoder_list[0].size()
            N = len(encoder_list)
            # (B,L,N*H)
            encoder_list = torch.cat(encoder_list,dim=2)
            encoder_list = encoder_list.view(B, L, N, H)
            encoder_out = torch.unsqueeze(encoder_list, dim=2)
            return {
                'encoder_out': encoder_out,  # T x B x C
                'encoder_padding_mask': encoder_padding_mask,  # B x T
            }
        else:
            return {
                'encoder_out': x,  # T x B x C
                'encoder_padding_mask': encoder_padding_mask,  # B x T
            }
    def max_positions(self):
        """Maximum input length supported by the encoder."""
        return self.embed_positions.max_positions()

    def upgrade_state_dict(self, state_dict):
        if isinstance(self.embed_positions, SinusoidalPositionalEmbedding):
            if 'encoder.embed_positions.weights' in state_dict:
                del state_dict['encoder.embed_positions.weights']
            if 'encoder.embed_positions._float_tensor' not in state_dict:
                state_dict['encoder.embed_positions._float_tensor'] = torch.FloatTensor()
        return state_dict


class StackTransformerDecoder(FairseqIncrementalDecoder):
    """Transformer decoder."""

    def __init__(self, args, dictionary, embed_tokens, left_pad=False):
        super().__init__(dictionary)
        self.dropout = args.dropout
        self.share_input_output_embed = args.share_decoder_input_output_embed

        embed_dim = embed_tokens.embedding_dim
        padding_idx = embed_tokens.padding_idx

        self.embed_tokens = embed_tokens
        self.embed_scale = math.sqrt(embed_dim)
        self.embed_positions = PositionalEmbedding(
            1024, embed_dim, padding_idx,
            left_pad=left_pad,
            learned=args.decoder_learned_pos,
        )

        self.layers = nn.ModuleList([])
        self.layers.extend([
            TransformerDecoderLayer(args)
            for i in range(args.decoder_layers)
        ])

        if not self.share_input_output_embed:
            self.embed_out = nn.Parameter(torch.Tensor(len(dictionary), embed_dim))
            nn.init.normal_(self.embed_out, mean=0, std=embed_dim ** -0.5)

        self.outer_stack = CreateOuterLayerStack(args, False)
        self.inner_stack_list = nn.ModuleList([])
        for i in range(args.decoder_layers):
            self.inner_stack_list.append(CreateLayerStack(args, False,True))

        if args.encoder_outer_stack_type == "transparent":
            self.fusion_weight = Parameter(torch.Tensor(args.decoder_layers,1 + args.encoder_layers))
            nn.init.xavier_uniform_(self.fusion_weight)
            #self.fusion_weight.data.fill_(1.0)
            self.use_fusion = True
        else:
            self.register_parameter('fusion_weight', None)
            self.use_fusion = False
        print(self.fusion_weight)
    def forward(self, prev_output_tokens, encoder_out, incremental_state=None):
        for i in range(len(self.inner_stack_list)):
            self.inner_stack_list[i].clear()
        self.outer_stack.clear()

        encoder_fusion = encoder_out["encoder_out"]
        #compute weight
        if self.use_fusion:
            Fusion_weight = F.dropout(self.fusion_weight, self.dropout, self.training)
            Fusion_weight_softmax = F.softmax(Fusion_weight, dim=1)
            #print(Fusion_weight)
            Fusion_weight_softmax = torch.unsqueeze(torch.unsqueeze(Fusion_weight_softmax, dim=0), 0)
            Fusion_weight_softmax = torch.unsqueeze(Fusion_weight_softmax, dim=-1)
            self.encoder_out = encoder_fusion * Fusion_weight_softmax
        # embed positions
        positions = self.embed_positions(
            prev_output_tokens,
            incremental_state=incremental_state,
        )

        if incremental_state is not None:
            prev_output_tokens = prev_output_tokens[:, -1:]
            positions = positions[:, -1:]

        # embed tokens and positions
        x = self.embed_scale * self.embed_tokens(prev_output_tokens)
        x += positions
        x = F.dropout(x, p=self.dropout, training=self.training)
        # B x T x C -> T x B x C
        x = x.transpose(0, 1).contiguous()
        self.inner_stack_list[0].push(x,use_dropout=False)
        self.outer_stack.push(x, use_dropout=False)

        # decoder layers
        for i,layer in enumerate(self.layers):
            if self.use_fusion:
                self.inner_stack_list[i], attn = layer(
                    self.inner_stack_list[i],
                    self.encoder_out[:, :, i],
                    encoder_out['encoder_padding_mask'],
                    incremental_state,
                    self.use_fusion
                )
            else:
                self.inner_stack_list[i], attn = layer(
                    self.inner_stack_list[i],
                    encoder_fusion,
                    encoder_out['encoder_padding_mask'],
                    incremental_state,
                )
            inner_out = self.inner_stack_list[i].pop()
            self.outer_stack.push(inner_out,use_dropout=False)
            self.inner_stack_list[i].clear()
            if i != len(self.layers)-1:
                self.inner_stack_list[i+1].push(self.outer_stack.pop(),use_dropout=False)

        x = self.outer_stack.pop()
        # T x B x C -> B x T x C
        x = x.transpose(0, 1)

        # project back to size of vocabulary
        if self.share_input_output_embed:
            x = F.linear(x, self.embed_tokens.weight)
        else:
            x = F.linear(x, self.embed_out)

        return x, attn

    def reorder_encoder_out(self, encoder_out_dict, new_order):
        if encoder_out_dict['encoder_padding_mask'] is not None:
            encoder_out_dict['encoder_padding_mask'] = \
                encoder_out_dict['encoder_padding_mask'].index_select(0, new_order)
        return encoder_out_dict

    def max_positions(self):
        """Maximum output length supported by the decoder."""
        return self.embed_positions.max_positions()

    def upgrade_state_dict(self, state_dict):
        if isinstance(self.embed_positions, SinusoidalPositionalEmbedding):
            if 'decoder.embed_positions.weights' in state_dict:
                del state_dict['decoder.embed_positions.weights']
            if 'decoder.embed_positions._float_tensor' not in state_dict:
                state_dict['decoder.embed_positions._float_tensor'] = torch.FloatTensor()
        return state_dict


class TransformerEncoderLayer(nn.Module):
    """Encoder layer block.

    In the original paper each operation (multi-head attention or FFN) is
    postprocessed with: dropout -> add residual -> layernorm.
    In the tensor2tensor code they suggest that learning is more robust when
    preprocessing each layer with layernorm and postprocessing with:
    dropout -> add residual.
    We default to the approach in the paper, but the tensor2tensor approach can
    be enabled by setting `normalize_before=True`.
    """

    def __init__(self, args):
        super().__init__()
        self.embed_dim = args.encoder_embed_dim
        self.self_attn = MultiheadAttention(
            self.embed_dim, args.encoder_attention_heads,
            dropout=args.attention_dropout,
        )
        self.dropout = args.dropout
        self.relu_dropout = args.relu_dropout
        #self.normalize_before = False
        self.fc1 = Linear(self.embed_dim, args.encoder_ffn_embed_dim)
        self.fc2 = Linear(args.encoder_ffn_embed_dim, self.embed_dim)
        #self.layer_norms = nn.ModuleList([LayerNorm(self.embed_dim) for i in range(2)])

    def forward(self, stack, encoder_padding_mask):
        x = stack.pop()
        #print('one layer:{}'.format(x))
        #x = self.maybe_layer_norm(0, x, before=True)
        x, _ = self.self_attn(query=x, key=x, value=x, key_padding_mask=encoder_padding_mask)
        stack.push(x)
        #x = self.maybe_layer_norm(0, x, after=True)

        x = stack.pop()
        #x = self.maybe_layer_norm(1, x, before=True)
        x = F.relu(self.fc1(x))
        x = F.dropout(x, p=self.relu_dropout, training=self.training)
        x = self.fc2(x)
        stack.push(x)
        #x = self.maybe_layer_norm(1, x, after=True)

        return stack

    def maybe_layer_norm(self, i, x, before=False, after=False):
        assert before ^ after
        if after ^ self.normalize_before:
            return self.layer_norms[i](x)
        else:
            return x


class TransformerDecoderLayer(nn.Module):
    """Decoder layer block."""

    def __init__(self, args):
        super().__init__()
        self.embed_dim = args.decoder_embed_dim
        self.self_attn = MultiheadAttention(
            self.embed_dim, args.decoder_attention_heads,
            dropout=args.attention_dropout,
        )
        self.dropout = args.dropout
        self.relu_dropout = args.relu_dropout
        #self.normalize_before = False
        self.encoder_attn = MultiheadAttention(
            self.embed_dim, args.decoder_attention_heads,
            dropout=args.attention_dropout,
        )
        self.fc1 = Linear(self.embed_dim, args.decoder_ffn_embed_dim)
        self.fc2 = Linear(args.decoder_ffn_embed_dim, self.embed_dim)
        #self.layer_norms = nn.ModuleList([LayerNorm(self.embed_dim) for i in range(3)])

    def forward(self, stack, encoder_out, encoder_padding_mask, incremental_state, use_fusion=False):
        x = stack.pop()
        #print('decoder one layer:{}'.format(x))
        #x = self.maybe_layer_norm(0, x, before=True)
        x, _ = self.self_attn(
            query=x,
            key=x,
            value=x,
            mask_future_timesteps=True,
            incremental_state=incremental_state,
            need_weights=False,
        )
        stack.push(x)
        #x = self.maybe_layer_norm(0, x, after=True)

        x = stack.pop()
        #x = self.maybe_layer_norm(1, x, before=True)
        if use_fusion:
            encoder_out = torch.mean(encoder_out,dim=2)
        x, attn = self.encoder_attn(
            query=x,
            key=encoder_out,
            value=encoder_out,
            key_padding_mask=encoder_padding_mask,
            incremental_state=incremental_state,
            static_kv=True,
        )
        stack.push(x)
        #x = self.maybe_layer_norm(1, x, after=True)

        x = stack.pop()
        #x = self.maybe_layer_norm(2, x, before=True)
        x = F.relu(self.fc1(x))
        x = F.dropout(x, p=self.relu_dropout, training=self.training)
        x = self.fc2(x)
        stack.push(x)
        #x = self.maybe_layer_norm(2, x, after=True)

        return stack, attn

    def maybe_layer_norm(self, i, x, before=False, after=False):
        assert before ^ after
        if after ^ self.normalize_before:
            return self.layer_norms[i](x)
        else:
            return x


def Embedding(num_embeddings, embedding_dim, padding_idx):
    m = nn.Embedding(num_embeddings, embedding_dim, padding_idx=padding_idx)
    nn.init.normal_(m.weight, mean=0, std=embedding_dim ** -0.5)
    return m


def LayerNorm(embedding_dim):
    m = nn.LayerNorm(embedding_dim)
    return m


def Linear(in_features, out_features, bias=True):
    m = nn.Linear(in_features, out_features, bias)
    nn.init.xavier_uniform_(m.weight)
    nn.init.constant_(m.bias, 0.)
    return m


def PositionalEmbedding(num_embeddings, embedding_dim, padding_idx, left_pad, learned=False):
    if learned:
        m = LearnedPositionalEmbedding(num_embeddings, embedding_dim, padding_idx, left_pad)
        nn.init.normal_(m.weight, mean=0, std=embedding_dim ** -0.5)
        nn.init.constant_(m.weight[padding_idx], 0)
    else:
        m = SinusoidalPositionalEmbedding(embedding_dim, padding_idx, left_pad, init_size=num_embeddings)
    return m


@register_model_architecture('stack_transformer', 'base')
def base_architecture(args):
    args.encoder_embed_dim = getattr(args, 'encoder_embed_dim', 512)
    args.encoder_ffn_embed_dim = getattr(args, 'encoder_ffn_embed_dim', 2048)
    args.encoder_layers = getattr(args, 'encoder_layers', 6)
    args.encoder_attention_heads = getattr(args, 'encoder_attention_heads', 8)
    args.decoder_embed_dim = getattr(args, 'decoder_embed_dim', args.encoder_embed_dim)
    args.decoder_ffn_embed_dim = getattr(args, 'decoder_ffn_embed_dim', args.encoder_ffn_embed_dim)
    args.decoder_layers = getattr(args, 'decoder_layers', 6)
    args.decoder_attention_heads = getattr(args, 'decoder_attention_heads', 8)
    args.attention_dropout = getattr(args, 'attention_dropout', 0.)
    args.relu_dropout = getattr(args, 'relu_dropout', 0.)
    args.dropout = getattr(args, 'dropout', 0.1)
    args.stack_dropout = getattr(args, 'stack_dropout', args.dropout)
    args.stack_relu_dropout = getattr(args, 'stack_relu_dropout', args.relu_dropout)
    args.stack_att_dropout = getattr(args, 'stack_att_dropout', args.attention_dropout)
    args.encoder_layer_normal_owned_type = getattr(args, 'encoder_layer_normal_owned_type', 'every_sublayer')
    args.decoder_layer_normal_owned_type = getattr(args, 'decoder_layer_normal_owned_type', 'every_sublayer')


@register_model_architecture('stack_transformer', 'residual_stack_transformer_iwslt_de_en')
def residual_stack_transformer_iwslt_de_en(args):
    args.encoder_embed_dim = getattr(args, 'encoder_embed_dim', 256)
    args.encoder_ffn_embed_dim = getattr(args, 'encoder_ffn_embed_dim', 512)
    args.encoder_attention_heads = getattr(args, 'encoder_attention_heads', 4)
    args.encoder_layers = getattr(args, 'encoder_layers', 3)
    args.decoder_embed_dim = getattr(args, 'decoder_embed_dim', 256)
    args.decoder_ffn_embed_dim = getattr(args, 'decoder_ffn_embed_dim', 512)
    args.decoder_attention_heads = getattr(args, 'decoder_attention_heads', 4)
    args.decoder_layers = getattr(args, 'decoder_layers', 3)
    base_architecture(args)

@register_model_architecture('stack_transformer', 'residual_stack_transformer_toy')
def residual_stack_transformer_toy(args):
    args.encoder_embed_dim = getattr(args, 'encoder_embed_dim', 64)
    args.encoder_ffn_embed_dim = getattr(args, 'encoder_ffn_embed_dim', 64)
    args.encoder_attention_heads = getattr(args, 'encoder_attention_heads', 4)
    args.encoder_layers = getattr(args, 'encoder_layers', 3)
    args.decoder_embed_dim = getattr(args, 'decoder_embed_dim', 64)
    args.decoder_ffn_embed_dim = getattr(args, 'decoder_ffn_embed_dim', 64)
    args.decoder_attention_heads = getattr(args, 'decoder_attention_heads', 4)
    args.decoder_layers = getattr(args, 'decoder_layers', 3)
    residual_stack_transformer_base_wmt_en_de(args)

@register_model_architecture('stack_transformer', 'residual_stack_transformer_base_wmt_en_de')
def residual_stack_transformer_base_wmt_en_de(args):
    args.encoder_inner_stack_type = 'residual'
    args.decoder_inner_stack_type = 'residual'
    args.encoder_outer_stack_type = 'None'
    args.decoder_outer_stack_type = 'None'
    #args.encoder_layers = 9
    base_architecture(args)

@register_model_architecture('stack_transformer', 'gru_stack_transformer_base_wmt_en_de')
def gru_stack_transformer_base_wmt_en_de(args):
    residual_stack_transformer_base_wmt_en_de(args)
    args.encoder_outer_stack_type = 'gru'
    args.decoder_outer_stack_type = 'gru'
    args.encoder_inner_stack_type = 'residual'
    args.decoder_inner_stack_type = 'residual'
    args.encoder_gru_owned_type = getattr(args, 'encoder_gru_owned_type', 'every_side')
    args.decoder_gru_owned_type = getattr(args, 'decoder_gru_owned_type', 'every_side')
    args.stack_hidden_size = getattr(args, 'stack_hidden_size', args.encoder_embed_dim)

@register_model_architecture('stack_transformer', 'aan_stack_transformer_base_wmt_en_de')
def aan_stack_transformer_base_wmt_en_de(args):
    residual_stack_transformer_base_wmt_en_de(args)
    args.encoder_stack_type = 'aan'
    args.decoder_stack_type = 'aan'
    args.add_layer_embed = True
    args.stack_hidden_size = getattr(args, 'stack_hidden_size', args.encoder_embed_dim)
    args.stack_inner_hidden_size = getattr(args, 'stack_inner_hidden_size', args.encoder_ffn_embed_dim)

@register_model_architecture('stack_transformer', 'mhsa_stack_transformer_base_wmt_en_de')
def mhsa_stack_transformer_base_wmt_en_de(args):
    residual_stack_transformer_base_wmt_en_de(args)
    args.encoder_inner_stack_type = 'mhsa'
    args.decoder_inner_stack_type = 'mhsa'
    args.encoder_outer_stack_type = 'dense'
    args.decoder_outer_stack_type = 'dense'
    args.encoder_layers = 15
    args.add_layer_embed = False
    args.stack_hidden_size = getattr(args, 'stack_hidden_size', args.encoder_embed_dim)
    args.stack_inner_hidden_size = getattr(args, 'stack_inner_hidden_size', args.encoder_ffn_embed_dim)
    args.stack_head_num = getattr(args, 'stack_head_num', 8)

@register_model_architecture('stack_transformer', 'dense_stack_transformer_base_wmt_en_de')
def dense_stack_transformer_base_wmt_en_de(args):
    residual_stack_transformer_base_wmt_en_de(args)
    args.encoder_inner_stack_type = 'residual'
    args.decoder_inner_stack_type = 'residual'
    args.encoder_outer_stack_type = 'dense'
    args.decoder_outer_stack_type = 'dense'
    args.encoder_layers = 22

@register_model_architecture('stack_transformer', 'transparent_stack_transformer_base_wmt_en_de')
def transparent_stack_transformer_base_wmt_en_de(args):
    residual_stack_transformer_base_wmt_en_de(args)
    args.encoder_inner_stack_type = 'residual'
    args.decoder_inner_stack_type = 'residual'
    args.encoder_outer_stack_type = 'transparent'
    args.decoder_outer_stack_type = 'None'
    args.encoder_layers = 6

@register_model_architecture('stack_transformer', 'transparent_stack_transformer_toy')
def transparent_stack_transformer_toy(args):
    residual_stack_transformer_toy(args)
    args.encoder_inner_stack_type = 'residual'
    args.decoder_inner_stack_type = 'residual'
    args.encoder_outer_stack_type = 'transparent'
    args.decoder_outer_stack_type = 'None'
    args.encoder_layers = 3

@register_model_architecture('stack_transformer', 'residual_stack_transformer_big')
def residual_stack_transformer_big(args):
    residual_stack_transformer_base_wmt_en_de(args)
    args.encoder_embed_dim = getattr(args, 'encoder_embed_dim', 1024)
    args.encoder_ffn_embed_dim = getattr(args, 'encoder_ffn_embed_dim', 4096)
    args.encoder_attention_heads = getattr(args, 'encoder_attention_heads', 16)
   # args.encoder_normalize_before = getattr(args, 'encoder_normalize_before', False)
    args.decoder_embed_dim = getattr(args, 'decoder_embed_dim', 1024)
    args.decoder_ffn_embed_dim = getattr(args, 'decoder_ffn_embed_dim', 4096)
    args.decoder_attention_heads = getattr(args, 'decoder_attention_heads', 16)
    args.dropout = getattr(args, 'dropout', 0.3) 
