# Source code for nlp_architect.models.intent_extraction

    # ******************************************************************************
    # Copyright 2017-2018 Intel Corporation
    #
    # Licensed under the Apache License, Version 2.0 (the "License");
    # you may not use this file except in compliance with the License.
    # You may obtain a copy of the License at
    #
    #     http://www.apache.org/licenses/LICENSE-2.0
    #
    # Unless required by applicable law or agreed to in writing, software
    # distributed under the License is distributed on an "AS IS" BASIS,
    # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    # See the License for the specific language governing permissions and
    # limitations under the License.
    # ******************************************************************************
    from __future__ import absolute_import, division, print_function, unicode_literals

    import tensorflow as tf

    from nlp_architect.nn.tensorflow.python.keras.layers.crf import CRF
    from nlp_architect.nn.tensorflow.python.keras.utils import load_model, save_model

    class IntentExtractionModel(object):
        """
        Intent Extraction model base class (using tf.keras)
        """

        def __init__(self):
            self.model = None

        def fit(self, x, y, epochs=1, batch_size=1, callbacks=None, validation=None):
            """
            Train a model given input samples and target labels.

            Args:
                x: input samples
                y: input sample labels
                epochs (:obj:`int`, optional): number of epochs to train
                batch_size (:obj:`int`, optional): batch size
                callbacks(:obj:`Callback`, optional): Keras compatible callbacks
                validation(:obj:`list` of :obj:`numpy.ndarray`, optional): optional validation data
                    to be evaluated when training
            """
            assert self.model, "Model was not initialized"
            self.model.fit(
                x,
                y,
                epochs=epochs,
                batch_size=batch_size,
                shuffle=True,
                validation_data=validation,
                callbacks=callbacks,
            )

        def predict(self, x, batch_size=1):
            """
            Get the prediction of the model on given input

            Args:
                x: samples to run through the model
                batch_size (:obj:`int`, optional): batch size:

            Returns:
                numpy.ndarray: predicted values by the model
            """
            assert self.model, "Model was not initialized"
            return self.model.predict(x, batch_size=batch_size)

        def save(self, path, exclude=None):
            """
            Save model to path

            Args:
                path (str): path to save model
                exclude (list, optional): a list of object fields to exclude when saving
            """
            assert self.model, "Model was not initialized"
            topology = {k: v for k, v in self.__dict__.items()}
            topology.pop("model")
            if exclude and isinstance(exclude, list):
                for x in exclude:
                    topology.pop(x)
            save_model(self.model, topology=topology, filepath=path)

        def load(self, path):
            """
            Load a trained model

            Args:
                path (str): path to model file
            """
            load_model(path, self)

        @property
        def input_shape(self):
            """:obj:`tuple`:Get input shape"""
            return self.model.layers[0].input_shape

        @staticmethod
        def _create_input_embed(sentence_len, is_extern_emb, token_emb_size, vocab_size):
            if is_extern_emb:
                in_layer = e_layer = tf.keras.layers.Input(
                    shape=(sentence_len, token_emb_size,), dtype="float32", name="tokens_input"
                )
            else:
                in_layer = tf.keras.layers.Input(
                    shape=(sentence_len,), dtype="int32", name="tokens_input"
                )
                e_layer = tf.keras.layers.Embedding(
                    vocab_size, token_emb_size, input_length=sentence_len, name="embedding_layer"
                )(in_layer)
            return in_layer, e_layer

        def load_embedding_weights(self, weights):
            """
            Load word embedding weights into the model embedding layer

            Args:
                weights (numpy.ndarray): 2D matrix of word weights
            """
            assert self.model is not None, (
                "Cannot assign weights, apply build() before trying to " "loading embedding weights "
            )
            emb_layer = self.model.get_layer(name="word_embedding")
            assert emb_layer.output_dim == weights.shape[1], "embedding vectors shape mismatch"
            emb_layer.set_weights([weights])

    class MultiTaskIntentModel(IntentExtractionModel):
        """
        Multi-Task Intent and Slot tagging model (using tf.keras)

        Args:
            use_cudnn (bool, optional): use GPU based model (CUDNNA cells)
        """

        def __init__(self, use_cudnn=False):
            super().__init__()
            self.model = None
            self.word_length = None
            self.num_labels = None
            self.num_intent_labels = None
            self.word_vocab_size = None
            self.char_vocab_size = None
            self.word_emb_dims = None
            self.char_emb_dims = None
            self.char_lstm_dims = None
            self.tagger_lstm_dims = None
            self.dropout = None
            self.use_cudnn = use_cudnn

        def build(
            self,
            word_length,
            num_labels,
            num_intent_labels,
            word_vocab_size,
            char_vocab_size,
            word_emb_dims=100,
            char_emb_dims=30,
            char_lstm_dims=30,
            tagger_lstm_dims=100,
            dropout=0.2,
        ):
            """
            Build a model

            Args:
                word_length (int): max word length (in characters)
                num_labels (int): number of slot labels
                num_intent_labels (int): number of intent classes
                word_vocab_size (int): word vocabulary size
                char_vocab_size (int): character vocabulary size
                word_emb_dims (int, optional): word embedding dimensions
                char_emb_dims (int, optional): character embedding dimensions
                char_lstm_dims (int, optional): character feature LSTM hidden size
                tagger_lstm_dims (int, optional): tagger LSTM hidden size
                dropout (float, optional): dropout rate
            """
            self.word_length = word_length
            self.num_labels = num_labels
            self.num_intent_labels = num_intent_labels
            self.word_vocab_size = word_vocab_size
            self.char_vocab_size = char_vocab_size
            self.word_emb_dims = word_emb_dims
            self.char_emb_dims = char_emb_dims
            self.char_lstm_dims = char_lstm_dims
            self.tagger_lstm_dims = tagger_lstm_dims
            self.dropout = dropout

            words_input = tf.keras.layers.Input(shape=(None,), name="words_input")
            embedding_layer = tf.keras.layers.Embedding(
                self.word_vocab_size, self.word_emb_dims, name="word_embedding"
            )
            word_embeddings = embedding_layer(words_input)
            word_embeddings = tf.keras.layers.Dropout(self.dropout)(word_embeddings)

            # create word character input and embeddings layer
            word_chars_input = tf.keras.layers.Input(
                shape=(None, self.word_length), name="word_chars_input"
            )
            char_embedding_layer = tf.keras.layers.Embedding(
                self.char_vocab_size,
                self.char_emb_dims,
                input_length=self.word_length,
                name="char_embedding",
            )
            # apply embedding to each word
            char_embeddings = char_embedding_layer(word_chars_input)
            # feed dense char vectors into BiLSTM
            char_embeddings = tf.keras.layers.TimeDistributed(
                tf.keras.layers.Bidirectional(self._rnn_cell(self.char_lstm_dims))
            )(char_embeddings)
            char_embeddings = tf.keras.layers.Dropout(self.dropout)(char_embeddings)

            # first BiLSTM layer (used for intent classification)
            first_bilstm_layer = tf.keras.layers.Bidirectional(
                self._rnn_cell(self.tagger_lstm_dims, return_sequences=True, return_state=True)
            )
            first_lstm_out = first_bilstm_layer(word_embeddings)

            lstm_y_sequence = first_lstm_out[:1][0]  # save y states of the LSTM layer
            states = first_lstm_out[1:]
            hf, _, hb, _ = states  # extract last hidden states
            h_state = tf.keras.layers.concatenate([hf, hb], axis=-1)
            intents = tf.keras.layers.Dense(
                self.num_intent_labels, activation="softmax", name="intent_classifier_output"
            )(h_state)

            # create the 2nd feature vectors
            combined_features = tf.keras.layers.concatenate([lstm_y_sequence, char_embeddings], axis=-1)

            # 2nd BiLSTM layer for label classification
            second_bilstm_layer = tf.keras.layers.Bidirectional(
                self._rnn_cell(self.tagger_lstm_dims, return_sequences=True)
            )(combined_features)
            second_bilstm_layer = tf.keras.layers.Dropout(self.dropout)(second_bilstm_layer)
            bilstm_out = tf.keras.layers.Dense(self.num_labels)(second_bilstm_layer)

            # feed BiLSTM vectors into CRF
            with tf.device("/cpu:0"):
                crf = CRF(self.num_labels, name="intent_slot_crf")
                labels = crf(bilstm_out)

            # compile the model
            model = tf.keras.Model(inputs=[words_input, word_chars_input], outputs=[intents, labels])

            # define losses and metrics
            loss_f = {
                "intent_classifier_output": "categorical_crossentropy",
                "intent_slot_crf": crf.loss,
            }
            metrics = {
                "intent_classifier_output": "categorical_accuracy",
                "intent_slot_crf": crf.viterbi_accuracy,
            }

            model.compile(loss=loss_f, optimizer=tf.train.AdamOptimizer(), metrics=metrics)
            self.model = model

        def _rnn_cell(self, units, **kwargs):
            if self.use_cudnn:
                rnn_cell = tf.keras.layers.CuDNNLSTM(units, **kwargs)
            else:
                rnn_cell = tf.keras.layers.LSTM(units, **kwargs)
            return rnn_cell

        # pylint: disable=arguments-differ
        def save(self, path):
            """
            Save model to path

            Args:
                path (str): path to save model
            """
            super().save(path, ["use_cudnn"])

    class Seq2SeqIntentModel(IntentExtractionModel):
        """
        Encoder Decoder Deep LSTM Tagger Model (using tf.keras)
        """

        def __init__(self):
            super().__init__()
            self.model = None
            self.vocab_size = None
            self.tag_labels = None
            self.token_emb_size = None
            self.encoder_depth = None
            self.decoder_depth = None
            self.lstm_hidden_size = None
            self.encoder_dropout = None
            self.decoder_dropout = None

        def build(
            self,
            vocab_size,
            tag_labels,
            token_emb_size=100,
            encoder_depth=1,
            decoder_depth=1,
            lstm_hidden_size=100,
            encoder_dropout=0.5,
            decoder_dropout=0.5,
        ):
            """
            Build the model

            Args:
                vocab_size (int): vocabulary size
                tag_labels (int): number of tag labels
                token_emb_size (int, optional): token embedding vector size
                encoder_depth (int, optional): number of encoder LSTM layers
                decoder_depth (int, optional): number of decoder LSTM layers
                lstm_hidden_size (int, optional): LSTM layers hidden size
                encoder_dropout (float, optional): encoder dropout
                decoder_dropout (float, optional): decoder dropout
            """
            self.vocab_size = vocab_size
            self.tag_labels = tag_labels
            self.token_emb_size = token_emb_size
            self.encoder_depth = encoder_depth
            self.decoder_depth = decoder_depth
            self.lstm_hidden_size = lstm_hidden_size
            self.encoder_dropout = encoder_dropout
            self.decoder_dropout = decoder_dropout

            words_input = tf.keras.layers.Input(shape=(None,), name="words_input")
            emb_layer = tf.keras.layers.Embedding(
                self.vocab_size, self.token_emb_size, name="word_embedding"
            )
            benc_in = emb_layer(words_input)

            assert self.encoder_depth > 0, "Encoder depth must be > 0"
            for i in range(self.encoder_depth):
                bencoder = tf.keras.layers.LSTM(
                    self.lstm_hidden_size,
                    return_sequences=True,
                    return_state=True,
                    go_backwards=True,
                    dropout=self.encoder_dropout,
                    name="encoder_blstm_{}".format(i),
                )(benc_in)
                benc_in = bencoder[0]
            b_states = bencoder[1:]
            benc_h, bene_c = b_states

            decoder_inputs = benc_in
            assert self.decoder_depth > 0, "Decoder depth must be > 0"
            for i in range(self.decoder_depth):
                decoder = tf.keras.layers.LSTM(
                    self.lstm_hidden_size, return_sequences=True, name="decoder_lstm_{}".format(i)
                )(decoder_inputs, initial_state=[benc_h, bene_c])
                decoder_inputs = decoder
            decoder_outputs = tf.keras.layers.Dropout(self.decoder_dropout)(decoder)
            decoder_predictions = tf.keras.layers.TimeDistributed(
                tf.keras.layers.Dense(self.tag_labels, activation="softmax"), name="decoder_classifier"
            )(decoder_outputs)

            self.model = tf.keras.Model(words_input, decoder_predictions)
            self.model.compile(
                optimizer=tf.train.AdamOptimizer(),
                loss="categorical_crossentropy",
                metrics=["categorical_accuracy"],
            )