import tensorflow as tf
import tensorflow_addons as tfa

import reader
from common import Common


class Model(tf.Module):
    def __init__(self, config, subtoken_vocab_size, target_vocab_size, nodes_vocab_size, target_to_index):
        super().__init__()
        self.config = config
        self.subtoken_vocab_shape = (subtoken_vocab_size, self.config.EMBEDDINGS_SIZE)
        self.target_vocab_shape = (target_vocab_size, self.config.EMBEDDINGS_SIZE)
        self.nodes_vocab_shape = (nodes_vocab_size, self.config.EMBEDDINGS_SIZE)
        self.target_to_index = target_to_index

        initializer = tf.initializers.VarianceScaling(scale=1.0,
                                                      mode='fan_out',
                                                      distribution='uniform')

        self.subtoken_vocab = tf.Variable(name='SUBTOKENS_VOCAB',
                                          shape=self.subtoken_vocab_shape,
                                          dtype=tf.float32,
                                          initial_value=initializer(self.subtoken_vocab_shape))

        self.target_words_vocab = tf.Variable(name='TARGET_WORDS_VOCAB',
                                              shape=self.target_vocab_shape,
                                              dtype=tf.float32,
                                              initial_value=initializer(self.target_vocab_shape))

        self.nodes_vocab = tf.Variable(name='NODES_VOCAB',
                                       shape=self.nodes_vocab_shape,
                                       dtype=tf.float32,
                                       initial_value=initializer(self.nodes_vocab_shape))

        self.rnn = None
        self.embed_dense_layer = None
        self.build_encoder()

        self.projection_layer = None
        self.attention_mechanism = None
        self.decoder_cell = None
        self.eval_decoder = None
        self.train_decoder = None
        self._beam_embedding = None
        self.build_decoder()

    def build_encoder(self):
        if self.config.BIRNN:
            rnn_cell_fw = tf.keras.layers.LSTMCell(self.config.RNN_SIZE // 2,
                                                   dropout=1 - self.config.RNN_DROPOUT_KEEP_PROB)
            rnn_cell_bw = tf.keras.layers.LSTMCell(self.config.RNN_SIZE // 2,
                                                   dropout=1 - self.config.RNN_DROPOUT_KEEP_PROB)
            self.rnn = tf.keras.layers.Bidirectional(layer=tf.keras.layers.RNN(rnn_cell_fw, return_state=True),
                                                     backward_layer=tf.keras.layers.RNN(rnn_cell_bw, go_backwards=True,
                                                                                        return_state=True),
                                                     merge_mode="concat",
                                                     dtype=tf.float32)
        else:
            rnn_cell = tf.keras.layers.LSTMCell(self.config.RNN_SIZE, dropout=1 - self.config.RNN_DROPOUT_KEEP_PROB)
            self.rnn = tf.keras.layers.RNN(rnn_cell, dtype=tf.float32, return_state=True)
        self.embed_dense_layer = tf.keras.layers.Dense(units=self.config.DECODER_SIZE,
                                                       activation=tf.nn.tanh, use_bias=False)

    def build_decoder(self):
        decoder_cells = [
            tf.keras.layers.LSTMCell(self.config.DECODER_SIZE, dropout=1 - self.config.RNN_DROPOUT_KEEP_PROB) for _ in
            range(self.config.NUM_DECODER_LAYERS)]
        self.decoder_cell = tf.keras.layers.StackedRNNCells(decoder_cells)
        self.projection_layer = tf.keras.layers.Dense(units=self.target_vocab_shape[0], use_bias=False)
        self.attention_mechanism = tfa.seq2seq.LuongAttention(units=self.config.DECODER_SIZE)
        should_save_alignment_history = self.config.BEAM_WIDTH == 0
        self.decoder_cell = tfa.seq2seq.AttentionWrapper(self.decoder_cell, self.attention_mechanism,
                                                         attention_layer_size=self.config.DECODER_SIZE,
                                                         alignment_history=should_save_alignment_history)
        if self.config.BEAM_WIDTH > 0:
            self.eval_decoder = tfa.seq2seq.BeamSearchDecoder(
                cell=self.decoder_cell,
                embedding_fn=lambda ids: tf.nn.embedding_lookup(self._beam_embedding, ids),
                beam_width=self.config.BEAM_WIDTH,
                output_layer=self.projection_layer,
                maximum_iterations=self.config.MAX_TARGET_PARTS + 1,
                length_penalty_weight=0.0)
        else:
            greedy_sampler = tfa.seq2seq.GreedyEmbeddingSampler()
            self.eval_decoder = tfa.seq2seq.BasicDecoder(cell=self.decoder_cell,
                                                         sampler=greedy_sampler,
                                                         maximum_iterations=self.config.MAX_TARGET_PARTS + 1,
                                                         output_layer=self.projection_layer)
        sampler = tfa.seq2seq.sampler.TrainingSampler()
        self.train_decoder = tfa.seq2seq.BasicDecoder(cell=self.decoder_cell,
                                                      sampler=sampler,
                                                      maximum_iterations=self.config.MAX_TARGET_PARTS + 1,
                                                      output_layer=self.projection_layer)

    @tf.function
    def run_encoder(self, input_tensors, is_training):
        path_source_indices = input_tensors[reader.PATH_SOURCE_INDICES_KEY]
        node_indices = input_tensors[reader.NODE_INDICES_KEY]
        path_target_indices = input_tensors[reader.PATH_TARGET_INDICES_KEY]
        valid_context_mask = input_tensors[reader.VALID_CONTEXT_MASK_KEY]
        path_source_lengths = input_tensors[reader.PATH_SOURCE_LENGTHS_KEY]
        path_lengths = input_tensors[reader.PATH_LENGTHS_KEY]
        path_target_lengths = input_tensors[reader.PATH_TARGET_LENGTHS_KEY]

        batched_contexts = self.compute_contexts(subtoken_vocab=self.subtoken_vocab,
                                                 nodes_vocab=self.nodes_vocab,
                                                 source_input=path_source_indices,
                                                 nodes_input=node_indices,
                                                 target_input=path_target_indices,
                                                 valid_mask=valid_context_mask,
                                                 path_source_lengths=path_source_lengths,
                                                 path_lengths=path_lengths,
                                                 path_target_lengths=path_target_lengths,
                                                 is_training=is_training)
        return batched_contexts

    def setup_attention_memory(self, batched_contexts):
        self.attention_mechanism.setup_memory(memory=batched_contexts)

    # @tf.function
    def run_decoder(self, batched_contexts, input_tensors, is_training):
        self.setup_attention_memory(batched_contexts)

        target_index = input_tensors[reader.TARGET_INDEX_KEY]
        valid_context_mask = input_tensors[reader.VALID_CONTEXT_MASK_KEY]
        batch_size = tf.shape(target_index)[0]
        outputs, final_states = self.decode_outputs(target_words_vocab=self.target_words_vocab,
                                                    target_input=target_index,
                                                    batch_size=batch_size,
                                                    batched_contexts=batched_contexts,
                                                    valid_mask=valid_context_mask,
                                                    is_training=is_training)

        return outputs, final_states

    def path_rnn_last_state(self, path_embed, path_lengths, valid_contexts_mask, is_training):
        # path_embed:           (batch, max_contexts, max_path_length+1, dim)
        # path_length:          (batch, max_contexts)
        # valid_contexts_mask:  (batch, max_contexts)
        max_contexts = tf.shape(path_embed)[1]

        # (batch * max_contexts, max_path_length+1, dim)
        flat_paths = tf.reshape(path_embed, shape=[-1, self.config.MAX_PATH_LENGTH,
                                                   self.config.EMBEDDINGS_SIZE])

        flat_valid_contexts_mask = tf.expand_dims(
            tf.sequence_mask(tf.reshape(path_lengths, [-1]), maxlen=self.config.MAX_PATH_LENGTH,
                             dtype=tf.float32), axis=-1)

        # https://github.com/tensorflow/tensorflow/issues/26974
        if self.config.BIRNN:
            res = self.rnn(inputs=flat_paths, mask=flat_valid_contexts_mask,
                           training=is_training)
            _, state_fw, _, state_bw, _ = res  # state = [mem, carry]
            final_rnn_state = tf.concat([state_fw, state_bw], axis=-1)  # (batch * max_contexts, rnn_size)
        else:
            _, state, _ = self.rnn(inputs=flat_paths, mask=flat_valid_contexts_mask, training=is_training)
            final_rnn_state = state

        return tf.reshape(final_rnn_state,
                          shape=[-1, max_contexts, self.config.RNN_SIZE])  # (batch, max_contexts, rnn_size)

    def compute_contexts(self, subtoken_vocab, nodes_vocab, source_input, nodes_input,
                         target_input, valid_mask, path_source_lengths, path_lengths, path_target_lengths, is_training):

        source_word_embed = tf.nn.embedding_lookup(params=subtoken_vocab,
                                                   ids=source_input)  # (batch, max_contexts, max_name_parts, dim)
        path_embed = tf.nn.embedding_lookup(params=nodes_vocab,
                                            ids=nodes_input)  # (batch, max_contexts, max_path_length+1, dim)
        target_word_embed = tf.nn.embedding_lookup(params=subtoken_vocab,
                                                   ids=target_input)  # (batch, max_contexts, max_name_parts, dim)

        source_word_mask = tf.expand_dims(
            tf.sequence_mask(path_source_lengths, maxlen=self.config.MAX_NAME_PARTS, dtype=tf.float32),
            -1)  # (batch, max_contexts, max_name_parts, 1)
        target_word_mask = tf.expand_dims(
            tf.sequence_mask(path_target_lengths, maxlen=self.config.MAX_NAME_PARTS, dtype=tf.float32),
            -1)  # (batch, max_contexts, max_name_parts, 1)

        source_words_sum = tf.reduce_sum(source_word_embed * source_word_mask,
                                         axis=2)  # (batch, max_contexts, dim)
        path_nodes_aggregation = self.path_rnn_last_state(path_embed, path_lengths,
                                                          valid_mask, is_training)  # (batch, max_contexts, rnn_size)
        target_words_sum = tf.reduce_sum(target_word_embed * target_word_mask, axis=2)  # (batch, max_contexts, dim)

        context_embed = tf.concat([source_words_sum, path_nodes_aggregation, target_words_sum],
                                  axis=-1)  # (batch, max_contexts, dim * 2 + rnn_size)
        if is_training:
            context_embed = tf.nn.dropout(context_embed, rate=1 - self.config.EMBEDDINGS_DROPOUT_KEEP_PROB)

        batched_embed = self.embed_dense_layer(inputs=context_embed)

        if not is_training and self.config.BEAM_WIDTH > 0:
            batched_embed = tfa.seq2seq.tile_batch(batched_embed, multiplier=self.config.BEAM_WIDTH)

        return batched_embed

    def decode_outputs(self, target_words_vocab, target_input, batch_size, batched_contexts, valid_mask, is_training):
        num_contexts_per_example = tf.math.count_nonzero(valid_mask, axis=-1)

        start_fill = tf.fill([batch_size],
                             self.target_to_index[Common.SOS])  # (batch, )

        contexts_sum = tf.reduce_sum(batched_contexts * tf.expand_dims(valid_mask, -1),
                                     axis=1)  # (batch_size, dim * 2 + rnn_size)
        contexts_average = tf.divide(contexts_sum, tf.cast(tf.expand_dims(num_contexts_per_example, -1), tf.float32))

        fake_encoder_state = tuple([contexts_average, contexts_average] for _ in
                                   range(self.config.NUM_DECODER_LAYERS))

        if not is_training:
            target_words_embedding = target_words_vocab
            if self.config.BEAM_WIDTH > 0:
                # https://medium.com/@dhirensk/tensorflow-addons-seq2seq-example-using-attention-and-beam-search-9f463b58bc6b
                decoder_initial_state = self.decoder_cell.get_initial_state(dtype=tf.float32,
                                                                            batch_size=batch_size * self.config.BEAM_WIDTH)
                decoder_initial_state = decoder_initial_state.clone(
                    cell_state=tfa.seq2seq.tile_batch(fake_encoder_state, multiplier=self.config.BEAM_WIDTH))
            else:
                decoder_initial_state = self.decoder_cell.get_initial_state(batch_size=batch_size, dtype=tf.float32)
                decoder_initial_state = decoder_initial_state.clone(cell_state=fake_encoder_state)

        else:
            # (batch, max_target_parts, dim * 2 + rnn_size)
            target_words_embedding = tf.nn.embedding_lookup(target_words_vocab,
                                                            tf.concat([tf.expand_dims(start_fill, -1), target_input],
                                                                      axis=-1))

            decoder_initial_state = self.decoder_cell.get_initial_state(batch_size=batch_size,
                                                                        dtype=tf.float32)
            decoder_initial_state = decoder_initial_state.clone(cell_state=fake_encoder_state)

        if is_training:
            outputs, final_states, final_sequence_lengths = self.train_decoder(
                target_words_embedding,
                training=True,
                initial_state=decoder_initial_state,
                sequence_length=tf.ones([batch_size], dtype=tf.int32) * (self.config.MAX_TARGET_PARTS + 1))
        else:
            if self.config.BEAM_WIDTH > 0:
                self._beam_embedding = target_words_embedding
                outputs, final_states, final_sequence_lengths = self.eval_decoder(
                    inputs=None,
                    training=False,
                    initial_state=decoder_initial_state,
                    start_tokens=start_fill,
                    end_token=self.target_to_index[Common.PAD])
            else:
                outputs, final_states, final_sequence_lengths = self.eval_decoder(
                    target_words_embedding,
                    training=False,
                    initial_state=decoder_initial_state,
                    start_tokens=start_fill,
                    end_token=0)

        return outputs, final_states

import tensorflow as tf
import tensorflow_addons as tfa

import reader
from common import Common


class Model(tf.Module):
    def __init__(self, config, subtoken_vocab_size, target_vocab_size, nodes_vocab_size, target_to_index):
        super().__init__()
        self.config = config
        self.subtoken_vocab_shape = (subtoken_vocab_size, self.config.EMBEDDINGS_SIZE)
        self.target_vocab_shape = (target_vocab_size, self.config.EMBEDDINGS_SIZE)
        self.nodes_vocab_shape = (nodes_vocab_size, self.config.EMBEDDINGS_SIZE)
        self.target_to_index = target_to_index

        initializer = tf.initializers.VarianceScaling(scale=1.0,
                                                      mode='fan_out',
                                                      distribution='uniform')

        self.subtoken_vocab = tf.Variable(name='SUBTOKENS_VOCAB',
                                          shape=self.subtoken_vocab_shape,
                                          dtype=tf.float32,
                                          initial_value=initializer(self.subtoken_vocab_shape))

        self.target_words_vocab = tf.Variable(name='TARGET_WORDS_VOCAB',
                                              shape=self.target_vocab_shape,
                                              dtype=tf.float32,
                                              initial_value=initializer(self.target_vocab_shape))

        self.nodes_vocab = tf.Variable(name='NODES_VOCAB',
                                       shape=self.nodes_vocab_shape,
                                       dtype=tf.float32,
                                       initial_value=initializer(self.nodes_vocab_shape))

        self.rnn = None
        self.embed_dense_layer = None
        self.build_encoder()

        self.projection_layer = None
        self.attention_mechanism = None
        self.decoder_cell = None
        self.eval_decoder = None
        self.train_decoder = None
        self._beam_embedding = None
        self.build_decoder()

    def build_encoder(self):
        if self.config.BIRNN:
            rnn_cell_fw = tf.keras.layers.LSTMCell(self.config.RNN_SIZE // 2,
                                                   dropout=1 - self.config.RNN_DROPOUT_KEEP_PROB)
            rnn_cell_bw = tf.keras.layers.LSTMCell(self.config.RNN_SIZE // 2,
                                                   dropout=1 - self.config.RNN_DROPOUT_KEEP_PROB)
            self.rnn = tf.keras.layers.Bidirectional(layer=tf.keras.layers.RNN(rnn_cell_fw, return_state=True),
                                                     backward_layer=tf.keras.layers.RNN(rnn_cell_bw, go_backwards=True,
                                                                                        return_state=True),
                                                     merge_mode="concat",
                                                     dtype=tf.float32)
        else:
            rnn_cell = tf.keras.layers.LSTMCell(self.config.RNN_SIZE, dropout=1 - self.config.RNN_DROPOUT_KEEP_PROB)
            self.rnn = tf.keras.layers.RNN(rnn_cell, dtype=tf.float32, return_state=True)
        self.embed_dense_layer = tf.keras.layers.Dense(units=self.config.DECODER_SIZE,
                                                       activation=tf.nn.tanh, use_bias=False)

    def build_decoder(self):
        decoder_cells = [
            tf.keras.layers.LSTMCell(self.config.DECODER_SIZE, dropout=1 - self.config.RNN_DROPOUT_KEEP_PROB) for _ in
            range(self.config.NUM_DECODER_LAYERS)]
        self.decoder_cell = tf.keras.layers.StackedRNNCells(decoder_cells)
        self.projection_layer = tf.keras.layers.Dense(units=self.target_vocab_shape[0], use_bias=False)
        self.attention_mechanism = tfa.seq2seq.LuongAttention(units=self.config.DECODER_SIZE)
        should_save_alignment_history = self.config.BEAM_WIDTH == 0
        self.decoder_cell = tfa.seq2seq.AttentionWrapper(self.decoder_cell, self.attention_mechanism,
                                                         attention_layer_size=self.config.DECODER_SIZE,
                                                         alignment_history=should_save_alignment_history)
        if self.config.BEAM_WIDTH > 0:
            self.eval_decoder = tfa.seq2seq.BeamSearchDecoder(
                cell=self.decoder_cell,
                embedding_fn=lambda ids: tf.nn.embedding_lookup(self._beam_embedding, ids),
                beam_width=self.config.BEAM_WIDTH,
                output_layer=self.projection_layer,
                maximum_iterations=self.config.MAX_TARGET_PARTS + 1,
                length_penalty_weight=0.0)
        else:
            greedy_sampler = tfa.seq2seq.GreedyEmbeddingSampler()
            self.eval_decoder = tfa.seq2seq.BasicDecoder(cell=self.decoder_cell,
                                                         sampler=greedy_sampler,
                                                         maximum_iterations=self.config.MAX_TARGET_PARTS + 1,
                                                         output_layer=self.projection_layer)
        sampler = tfa.seq2seq.sampler.TrainingSampler()
        self.train_decoder = tfa.seq2seq.BasicDecoder(cell=self.decoder_cell,
                                                      sampler=sampler,
                                                      maximum_iterations=self.config.MAX_TARGET_PARTS + 1,
                                                      output_layer=self.projection_layer)

    @tf.function
    def run_encoder(self, input_tensors, is_training):
        path_source_indices = input_tensors[reader.PATH_SOURCE_INDICES_KEY]
        node_indices = input_tensors[reader.NODE_INDICES_KEY]
        path_target_indices = input_tensors[reader.PATH_TARGET_INDICES_KEY]
        valid_context_mask = input_tensors[reader.VALID_CONTEXT_MASK_KEY]
        path_source_lengths = input_tensors[reader.PATH_SOURCE_LENGTHS_KEY]
        path_lengths = input_tensors[reader.PATH_LENGTHS_KEY]
        path_target_lengths = input_tensors[reader.PATH_TARGET_LENGTHS_KEY]

        batched_contexts = self.compute_contexts(subtoken_vocab=self.subtoken_vocab,
                                                 nodes_vocab=self.nodes_vocab,
                                                 source_input=path_source_indices,
                                                 nodes_input=node_indices,
                                                 target_input=path_target_indices,
                                                 valid_mask=valid_context_mask,
                                                 path_source_lengths=path_source_lengths,
                                                 path_lengths=path_lengths,
                                                 path_target_lengths=path_target_lengths,
                                                 is_training=is_training)
        return batched_contexts

    def setup_attention_memory(self, batched_contexts):
        self.attention_mechanism.setup_memory(memory=batched_contexts)

    # @tf.function
    def run_decoder(self, batched_contexts, input_tensors, is_training):
        self.setup_attention_memory(batched_contexts)

        target_index = input_tensors[reader.TARGET_INDEX_KEY]
        valid_context_mask = input_tensors[reader.VALID_CONTEXT_MASK_KEY]
        batch_size = tf.shape(target_index)[0]
        outputs, final_states = self.decode_outputs(target_words_vocab=self.target_words_vocab,
                                                    target_input=target_index,
                                                    batch_size=batch_size,
                                                    batched_contexts=batched_contexts,
                                                    valid_mask=valid_context_mask,
                                                    is_training=is_training)

        return outputs, final_states

    def path_rnn_last_state(self, path_embed, path_lengths, valid_contexts_mask, is_training):
        # path_embed:           (batch, max_contexts, max_path_length+1, dim)
        # path_length:          (batch, max_contexts)
        # valid_contexts_mask:  (batch, max_contexts)
        max_contexts = tf.shape(path_embed)[1]

        # (batch * max_contexts, max_path_length+1, dim)
        flat_paths = tf.reshape(path_embed, shape=[-1, self.config.MAX_PATH_LENGTH,
                                                   self.config.EMBEDDINGS_SIZE])

        flat_valid_contexts_mask = tf.expand_dims(
            tf.sequence_mask(tf.reshape(path_lengths, [-1]), maxlen=self.config.MAX_PATH_LENGTH,
                             dtype=tf.float32), axis=-1)

        # https://github.com/tensorflow/tensorflow/issues/26974
        if self.config.BIRNN:
            res = self.rnn(inputs=flat_paths, mask=flat_valid_contexts_mask,
                           training=is_training)
            _, state_fw, _, state_bw, _ = res  # state = [mem, carry]
            final_rnn_state = tf.concat([state_fw, state_bw], axis=-1)  # (batch * max_contexts, rnn_size)
        else:
            _, state, _ = self.rnn(inputs=flat_paths, mask=flat_valid_contexts_mask, training=is_training)
            final_rnn_state = state

        return tf.reshape(final_rnn_state,
                          shape=[-1, max_contexts, self.config.RNN_SIZE])  # (batch, max_contexts, rnn_size)

    def compute_contexts(self, subtoken_vocab, nodes_vocab, source_input, nodes_input,
                         target_input, valid_mask, path_source_lengths, path_lengths, path_target_lengths, is_training):

        source_word_embed = tf.nn.embedding_lookup(params=subtoken_vocab,
                                                   ids=source_input)  # (batch, max_contexts, max_name_parts, dim)
        path_embed = tf.nn.embedding_lookup(params=nodes_vocab,
                                            ids=nodes_input)  # (batch, max_contexts, max_path_length+1, dim)
        target_word_embed = tf.nn.embedding_lookup(params=subtoken_vocab,
                                                   ids=target_input)  # (batch, max_contexts, max_name_parts, dim)

        source_word_mask = tf.expand_dims(
            tf.sequence_mask(path_source_lengths, maxlen=self.config.MAX_NAME_PARTS, dtype=tf.float32),
            -1)  # (batch, max_contexts, max_name_parts, 1)
        target_word_mask = tf.expand_dims(
            tf.sequence_mask(path_target_lengths, maxlen=self.config.MAX_NAME_PARTS, dtype=tf.float32),
            -1)  # (batch, max_contexts, max_name_parts, 1)

        source_words_sum = tf.reduce_sum(source_word_embed * source_word_mask,
                                         axis=2)  # (batch, max_contexts, dim)
        path_nodes_aggregation = self.path_rnn_last_state(path_embed, path_lengths,
                                                          valid_mask, is_training)  # (batch, max_contexts, rnn_size)
        target_words_sum = tf.reduce_sum(target_word_embed * target_word_mask, axis=2)  # (batch, max_contexts, dim)

        context_embed = tf.concat([source_words_sum, path_nodes_aggregation, target_words_sum],
                                  axis=-1)  # (batch, max_contexts, dim * 2 + rnn_size)
        if is_training:
            context_embed = tf.nn.dropout(context_embed, rate=1 - self.config.EMBEDDINGS_DROPOUT_KEEP_PROB)

        batched_embed = self.embed_dense_layer(inputs=context_embed)

        if not is_training and self.config.BEAM_WIDTH > 0:
            batched_embed = tfa.seq2seq.tile_batch(batched_embed, multiplier=self.config.BEAM_WIDTH)

        return batched_embed

    def decode_outputs(self, target_words_vocab, target_input, batch_size, batched_contexts, valid_mask, is_training):
        num_contexts_per_example = tf.math.count_nonzero(valid_mask, axis=-1)

        start_fill = tf.fill([batch_size],
                             self.target_to_index[Common.SOS])  # (batch, )

        contexts_sum = tf.reduce_sum(batched_contexts * tf.expand_dims(valid_mask, -1),
                                     axis=1)  # (batch_size, dim * 2 + rnn_size)
        contexts_average = tf.divide(contexts_sum, tf.cast(tf.expand_dims(num_contexts_per_example, -1), tf.float32))

        fake_encoder_state = tuple([contexts_average, contexts_average] for _ in
                                   range(self.config.NUM_DECODER_LAYERS))

        if not is_training:
            target_words_embedding = target_words_vocab
            if self.config.BEAM_WIDTH > 0:
                # https://medium.com/@dhirensk/tensorflow-addons-seq2seq-example-using-attention-and-beam-search-9f463b58bc6b
                decoder_initial_state = self.decoder_cell.get_initial_state(dtype=tf.float32,
                                                                            batch_size=batch_size * self.config.BEAM_WIDTH)
                decoder_initial_state = decoder_initial_state.clone(
                    cell_state=tfa.seq2seq.tile_batch(fake_encoder_state, multiplier=self.config.BEAM_WIDTH))
            else:
                decoder_initial_state = self.decoder_cell.get_initial_state(batch_size=batch_size, dtype=tf.float32)
                decoder_initial_state = decoder_initial_state.clone(cell_state=fake_encoder_state)

        else:
            # (batch, max_target_parts, dim * 2 + rnn_size)
            target_words_embedding = tf.nn.embedding_lookup(target_words_vocab,
                                                            tf.concat([tf.expand_dims(start_fill, -1), target_input],
                                                                      axis=-1))

            decoder_initial_state = self.decoder_cell.get_initial_state(batch_size=batch_size,
                                                                        dtype=tf.float32)
            decoder_initial_state = decoder_initial_state.clone(cell_state=fake_encoder_state)

        if is_training:
            outputs, final_states, final_sequence_lengths = self.train_decoder(
                target_words_embedding,
                training=True,
                initial_state=decoder_initial_state,
                sequence_length=tf.ones([batch_size], dtype=tf.int32) * (self.config.MAX_TARGET_PARTS + 1))
        else:
            if self.config.BEAM_WIDTH > 0:
                self._beam_embedding = target_words_embedding
                outputs, final_states, final_sequence_lengths = self.eval_decoder(
                    inputs=None,
                    training=False,
                    initial_state=decoder_initial_state,
                    start_tokens=start_fill,
                    end_token=self.target_to_index[Common.PAD])
            else:
                outputs, final_states, final_sequence_lengths = self.eval_decoder(
                    target_words_embedding,
                    training=False,
                    initial_state=decoder_initial_state,
                    start_tokens=start_fill,
                    end_token=0)

        return outputs, final_states

# @tf.function # path_embed:           (batch, max_contexts, max_path_length+1, dim) # path_length:          (batch, max_contexts) # valid_contexts_mask:  (batch, max_contexts) # (batch * max_contexts, max_path_length+1, dim) # https://github.com/tensorflow/tensorflow/issues/26974 # state = [mem, carry] # (batch * max_contexts, rnn_size) # (batch, max_contexts, rnn_size) # (batch, max_contexts, max_name_parts, dim) # (batch, max_contexts, max_path_length+1, dim) # (batch, max_contexts, max_name_parts, dim) # (batch, max_contexts, max_name_parts, 1) # (batch, max_contexts, max_name_parts, 1) # (batch, max_contexts, dim) # (batch, max_contexts, rnn_size) # (batch, max_contexts, dim) # (batch, max_contexts, dim * 2 + rnn_size) # (batch, ) # (batch_size, dim * 2 + rnn_size) # https://medium.com/@dhirensk/tensorflow-addons-seq2seq-example-using-attention-and-beam-search-9f463b58bc6b # (batch, max_target_parts, dim * 2 + rnn_size)