import tensorflow as tf
import numpy as np
from deep_models.LSTM import LSTM
from datetime import datetime
from sklearn.exceptions import NotFittedError

he_init = tf.contrib.layers.variance_scaling_initializer()


class HARNN(LSTM):

    def __init__(self, h_max_length=20, b_max_length=200, trainable=False, lstm_layers=2, mlp_layers=1,
                 num_neurons=[128, 128, 32], share_parameters=True, average_pooling=False,
                 optimizer=tf.train.AdamOptimizer,
                 learning_rate=0.001, batch_size=128, activation=tf.nn.relu, initializer=he_init, num_epoch=40,
                 batch_norm_momentum=None, dropout_rate=None,
                 max_check_without_progress=20, show_progress=10, tensorboard_logdir=None, random_state=None,
                 embedding=None, l2_lambda=0.01, word_output_size=64, sent_output_size=64, vocab_size=None):
        LSTM.__init__(self, h_max_length, b_max_length, trainable, lstm_layers, mlp_layers, num_neurons,
                      share_parameters, average_pooling, optimizer, learning_rate, batch_size, activation, initializer,
                      num_epoch, batch_norm_momentum, dropout_rate,
                      max_check_without_progress, show_progress, tensorboard_logdir, random_state, embedding, l2_lambda,
                      vocab_size)

        self.word_output_size = word_output_size
        self.sent_output_size = sent_output_size
        self.vocab_size = vocab_size
        self.embedding_size = 100
        if self.embedding is None and self.vocab_size is None:
            raise Exception("Either embedding or vocab_size must be setted!")

    def gru_cell(self, num_neuron):

        gru = tf.contrib.rnn.GRUCell(num_neuron)
        if self.dropout_rate:
            gru = tf.contrib.rnn.DropoutWrapper(gru, input_keep_prob=self.keep_prob, output_keep_prob=self.keep_prob)
        return gru

    def _bidirectional_rnn(self, inputs, inputs_length, scope=None):

        with tf.variable_scope(scope or 'birnn') as scope:

            if self.lstm_layers == 1:
                rnn_cells = self.gru_cell(self.num_neurons[0])
            elif self.lstm_layers > 1:
                rnn_cells = tf.nn.rnn_cell.MultiRNNCell(
                    [self.gru_cell(i) for i in self.num_neurons[:self.lstm_layers]])

            ((fw_outputs, bw_outputs), (fw_states, bw_states)) = tf.nn.bidirectional_dynamic_rnn(rnn_cells, rnn_cells,
                                                                                                 inputs, inputs_length,
                                                                                                 dtype=tf.float32)
            outputs = tf.concat([fw_outputs, bw_outputs], axis=2)

            if self.lstm_layers > 1:
                fw_states = fw_states[self.lstm_layers - 1]
                bw_states = bw_states[self.lstm_layers - 1]

        return outputs, fw_states, bw_states

    def _attention(self, inputs, output_size, initializer=tf.contrib.layers.xavier_initializer(), activation_fn=tf.tanh,
                   scope=None):

        with tf.variable_scope(scope or 'attention') as scope:
            attetion_context_vector = tf.get_variable(name='attention_context_vector', shape=[output_size],
                                                      initializer=initializer, dtype=tf.float32)
            input_projection = tf.contrib.layers.fully_connected(inputs, output_size, activation_fn=activation_fn,
                                                                 scope=scope)

            vector_attn = tf.reduce_sum(tf.multiply(input_projection, attetion_context_vector), axis=2, keep_dims=True)
            attention_weights = tf.nn.softmax(vector_attn, axis=1)
            weighted_projection = tf.multiply(input_projection, attention_weights)

            outputs = tf.reduce_mean(weighted_projection, axis=1)

            return outputs

    def _hierarichcal_attention(self, inputs, sizes, sent_sizes, scope):

        (doc_size, sent_size, word_size, embed_size) = tf.unstack(tf.shape(inputs))
        print(sizes.get_shape())

        with tf.variable_scope(scope):
            word_level_inputs = tf.reshape(inputs, [doc_size * sent_size, word_size, self.embedding_size])
            word_level_lengths = tf.reshape(sent_sizes, [doc_size * sent_size])

            with tf.variable_scope("word") as scope:
                word_encoder_output, _, _ = self._bidirectional_rnn(word_level_inputs, word_level_lengths, scope)

                with tf.variable_scope("attention") as scope:
                    word_level_output = self._attention(word_encoder_output, self.word_output_size, scope=scope)

                    with tf.variable_scope("dropout") as scope:
                        word_level_output = tf.layers.dropout(word_level_output, rate=self.dropout_rate,
                                                              training=self._training)

            sentence_level_inputs = tf.reshape(word_level_output, (doc_size, sent_size, self.word_output_size))
            with tf.variable_scope("sentence") as scope:
                sentence_encoder_outputs, _, _ = self._bidirectional_rnn(sentence_level_inputs, sizes, scope=scope)

                with tf.variable_scope("attention") as scope:
                    sent_level_output = self._attention(sentence_encoder_outputs, self.sent_output_size, scope=scope)

                    with tf.variable_scope("dropout") as scope:
                        sent_level_output = tf.layers.dropout(sent_level_output, rate=self.dropout_rate)

        return sent_level_output

    def _add_embedding(self, inputs, scope):

        with tf.variable_scope(scope):
            with tf.variable_scope("embedding_lookup") as scope:
                if self.embedding is None and self.vocab_size is not None:
                    embedding = tf.get_variable(
                        name="embedding_matrix",
                        shape=[self.vocab_size, self.embedding_size],
                        initializer=tf.contrib.layers.xavier_initializer(),
                        dtype=tf.float32)
                # else:
                #     embedding = tf.Variable(self.embedding,trainable=self.trainable)
                inputs_embedded = tf.nn.embedding_lookup(embedding, inputs)
                return inputs_embedded

    def _ann(self, head_inputs, body_inputs, h_sizes, b_sizes, h_sent_sizes, b_sent_sizes):

        with tf.variable_scope("embedding_lookup") as scope:
            heads_embeddings = self._add_embedding(head_inputs, scope)
            scope.reuse_variables()
            body_embeddings = self._add_embedding(body_inputs, scope)

        if self.share_parameters:
            with tf.variable_scope("heads_and_bodies") as scope:
                head_sent_level_outputs = self._hierarichcal_attention(heads_embeddings, h_sizes, h_sent_sizes, scope)
                scope.reuse_variables()
                body_sent_level_outputs = self._hierarichcal_attention(body_embeddings, b_sizes, b_sent_sizes, scope)
        else:
            with tf.variable_scope("heads") as scope:
                head_sent_level_outputs = self._hierarichcal_attention(heads_embeddings, h_sizes, h_sent_sizes, scope)
            with tf.variable_scope("bodies") as scope:
                body_sent_level_outputs = self._hierarichcal_attention(body_embeddings, b_sizes, b_sent_sizes, scope)

        output = tf.concat([head_sent_level_outputs, body_sent_level_outputs], axis=1)

        for layer in range(self.mlp_layers):

            if self.dropout_rate:
                output = tf.layers.dropout(output, rate=self.dropout_rate, training=self._training)
            output = tf.layers.dense(output, self.num_neurons[self.lstm_layers + layer], activation=self.activation,
                                     kernel_initializer=self.initializer,
                                     name="hidden{}".format(layer + 1))

            output = self.activation(output, name="hidden{}_out".format(layer + 1))

        return output

    def _construct_graph(self, n_outputs, weights):

        if self.randome_state:
            tf.set_random_seed(self.randome_state)
            np.random.seed(self.randome_state)
        X_heads = tf.placeholder(dtype=tf.int32, shape=[None, None, None], name="X_heads")
        X_bodies = tf.placeholder(dtype=tf.int32, shape=[None, None, None], name="X_bodies")
        X_head_sizes = tf.placeholder(dtype=tf.int32, shape=[None], name="head_sizes")
        X_body_sizes = tf.placeholder(dtype=tf.int32, shape=[None], name="body_sizes")
        X_head_sent_sizes = tf.placeholder(dtype=tf.int32, shape=[None, None], name="head_sent_sizes")
        X_body_sent_sizes = tf.placeholder(dtype=tf.int32, shape=[None, None], name="body_sent_sizes")

        # X_addtional = tf.placeholder(tf.float32,shape=[None,features_length],name="additonal_features")
        y_ = tf.placeholder(tf.int32, shape=[None], name="y")
        weights = tf.constant(weights, tf.float32)

        if self.dropout_rate:
            self._training = tf.placeholder_with_default(False, shape=[], name="training")
            self.keep_prob = tf.cond(self._training, lambda: tf.constant(1 - self.dropout_rate),
                                     lambda: tf.constant(1.0))
        else:
            self._training = None

        pre_output = self._ann(X_heads, X_bodies, X_head_sizes, X_body_sizes, X_head_sent_sizes, X_body_sent_sizes)
        logits = tf.layers.dense(pre_output, n_outputs, kernel_initializer=he_init, name="logits")
        # weighted_logits = tf.multiply(logits,weights)
        probabilities = tf.nn.softmax(logits, name="probabilities")

        xentropy = tf.nn.sparse_softmax_cross_entropy_with_logits(labels=y_, logits=logits)
        loss = tf.reduce_mean(xentropy, name="loss")
        variables = tf.trainable_variables()
        for v in variables:
            print(v.name)
        l2_loss = tf.add_n(
            [tf.nn.l2_loss(v) for v in variables if 'bias' not in v.name and 'Variable' not in v.name]) * self.l2_lambda
        loss += l2_loss

        optimizer = self.optimizer(learning_rate=self.learning_rate)
        update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
        with tf.control_dependencies(update_ops):
            training_op = optimizer.minimize(loss)

        correct = tf.nn.in_top_k(logits, y_, 1)
        accuracy = tf.reduce_mean(tf.cast(correct, tf.float32), name="accuracy")

        init = tf.global_variables_initializer()
        saver = tf.train.Saver()

        if self.tensorboard_logdir:
            now = datetime.utcnow().strftime('%Y%m%d-%H%M%S')
            tb_logdir = self.tensorboard_logdir + "/run{}".format(now)
            cost_summary = tf.summary.scalar("validation_loss", loss)
            acc_summary = tf.summary.scalar("validation_accuracy", accuracy)
            merged_summary = tf.summary.merge_all()
            file_writer = tf.summary.FileWriter(tb_logdir, tf.get_default_graph())

            self._merged_summary = merged_summary
            self._file_writer = file_writer

        self._X_head, self._X_body, self._X_h_sizes, self._X_b_sizes, self._X_h_sent_sizes, self._X_b_sent_sizes, self.y = X_heads, X_bodies, X_head_sizes, X_body_sizes, X_head_sent_sizes, X_body_sent_sizes, y_
        self._logits = logits
        self._probabilites = probabilities
        self._loss = loss
        self._training_op = training_op
        self._accuracy = accuracy
        self._init, self._saver = init, saver
        # self._X_additional = X_addtional

    def get_batch_attention(self, h_np, b_np, h_sizes, b_sizes, h_sent_sizes, b_sent_sizes, y):

        num_batches = h_np.shape[0] // self.batch_size
        batches = []

        for i in range(num_batches):
            if (i + 1) * self.batch_size > h_np.shape[0]:
                batch_h_np = h_np[i * self.batch_size:h_np.shape[0] - 1, :]
                batch_b_np = b_np[i * self.batch_size:h_np.shape[0] - 1, :]
                batch_h_sizes = h_sizes[i * self.batch_size:h_np.shape[0] - 1]
                batch_b_sizes = b_sizes[i * self.batch_size:h_np.shape[0] - 1]
                batch_h_sent_sizes = h_sent_sizes[i * self.batch_size:h_np.shape[0] - 1, :]
                batch_b_sent_sizes = b_sent_sizes[i * self.batch_size:h_np.shape[0] - 1, :]
                batch_y = y[i * self.batch_size:h_np.shape[0] - 1]
                # batch_additional = additional_features[i*self.batch_size:h_np.shape[0]-1]
            else:
                batch_h_np = h_np[i * self.batch_size:(i + 1) * self.batch_size, :]
                batch_b_np = b_np[i * self.batch_size:(i + 1) * self.batch_size, :]
                batch_h_sizes = h_sizes[i * self.batch_size:(i + 1) * self.batch_size]
                batch_b_sizes = b_sizes[i * self.batch_size:(i + 1) * self.batch_size]
                batch_h_sent_sizes = h_sent_sizes[i * self.batch_size:(i + 1) * self.batch_size, :]
                batch_b_sent_sizes = b_sent_sizes[i * self.batch_size:(i + 1) * self.batch_size, :]
                batch_y = y[i * self.batch_size:(i + 1) * self.batch_size]
                # batch_additional = additional_features[i*self.batch_size:(i+1)*self.batch_size]
            batches.append(
                (batch_h_np, batch_b_np, batch_h_sizes, batch_b_sizes, batch_h_sent_sizes, batch_b_sent_sizes, batch_y))

        return batches

    def sort_base_bodies(self, h_v, b_v, y):

        assert len(h_v) == len(b_v) == len(y)
        b_lengths = [len(b_vector) for b_vector in b_v]

        sorted_h = []
        sorted_b = []
        sorted_y = []
        for length in range(min(b_lengths), max(b_lengths)):
            for idx, vectors in enumerate(b_v):
                if len(vectors) == length:
                    sorted_h.append(h_v[idx])
                    sorted_b.append(b_v[idx])
                    sorted_y.append(y[idx])

        return sorted_h, sorted_b, sorted_y

    def get_batch(self, h_v, b_v, y=None):

        assert len(h_v) == len(b_v)

        for batch_i in range(0, len(b_v) // self.batch_size):
            start_i = batch_i * self.batch_size
            h_batch = h_v[start_i:start_i + self.batch_size]
            b_batch = b_v[start_i:start_i + self.batch_size]
            if y is not None:
                y_batch = y[start_i:start_i + self.batch_size]

            h_lengths_batch = [len(h) for h in h_batch]
            b_lengths_batch = [len(b) for b in b_batch]

            def padding(vectors, lengths):

                max_length = max(lengths)
                np_batch = np.zeros(shape=(len(vectors), max_length, 4096))
                for i, sents_vector in enumerate(vectors):
                    for j, sent in enumerate(sents_vector):
                        np_batch[i, j, :] = sent
                return np_batch

            np_h_batch = padding(h_batch, h_lengths_batch)
            np_b_batch = padding(b_batch, b_lengths_batch)
            h_lengths_batch = np.asarray(h_lengths_batch, np.int32)
            b_lengths_batch = np.asarray(b_lengths_batch, np.int32)

            if y is not None:
                yield np_h_batch, np_b_batch, h_lengths_batch, b_lengths_batch, y_batch
            else:
                yield np_h_batch, np_b_batch, h_lengths_batch, b_lengths_batch

    def fit(self, X, y, valid_X=None, y_valid=None):
        self.close_session()
        y = np.array(y)

        if len(y.shape) == 2:
            y = np.argmax(y, axis=1)

        self._classes, counts = np.unique(y, return_counts=True)
        weights = [count / sum(counts) for count in counts]
        weights = np.asarray(weights, np.float32)
        n_outputs = len(self._classes)

        self._graph = tf.Graph()
        h_np, b_np, h_sizes, b_sizes, h_sent_sizes, b_sent_sizes = X['h_np'], X['b_np'], X['h_sizes'], X['b_sizes'], X[
            'h_sent_sizes'], X['b_sent_sizes']
        #        features_length = addtional_features.shape[1]

        # print("head shape{}".format(h_sizes.shape))

        with self._graph.as_default():
            self._construct_graph(n_outputs, weights)

        checks_without_progress = 0
        best_acc = np.inf
        best_parameters = None
        gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.333)
        self._session = tf.Session(graph=self._graph, config=tf.ConfigProto(gpu_options=gpu_options))

        with self._session.as_default() as sess:
            self._init.run()
            num_instances = h_np.shape[0]
            for epoch in range(self.num_epoch):
                losses = []
                accs = []
                rnd_idx = np.random.permutation(num_instances)
                for rnd_indices in np.array_split(rnd_idx, num_instances // self.batch_size):

                    X_head_batch, X_body_batch, X_h_sizes_batch, X_b_sizes_batch, X_h_sent_sizes_batch, X_b_sent_sizes_batch, y_batch = \
                    h_np[rnd_indices], b_np[rnd_indices], \
                    h_sizes[rnd_indices], b_sizes[rnd_indices], \
                    h_sent_sizes[rnd_indices], b_sent_sizes[rnd_indices], \
                    y[rnd_indices]
                    # print("batch head shape{}".format(X_h_sizes_batch.shape))
                    y_batch = np.asarray(y_batch)
                    feed_dict = {self._X_head: X_head_batch, self._X_body: X_body_batch,
                                 self._X_h_sizes: X_h_sizes_batch, self._X_b_sizes: X_b_sizes_batch,
                                 self._X_h_sent_sizes: X_h_sent_sizes_batch, self._X_b_sent_sizes: X_b_sent_sizes_batch,
                                 self.y: y_batch}
                    if self._training is not None:
                        feed_dict[self._training] = True

                    train_acc, _, loss = sess.run([self._accuracy, self._training_op, self._loss], feed_dict=feed_dict)
                    losses.append(loss)
                    accs.append(train_acc)
                average_loss = sum(losses) / len(losses)
                average_acc = sum(accs) / len(accs)

                if valid_X is not None and y_valid is not None:

                    print("validation phase")
                    # valid_h_np, valid_b_np, valid_h_lenseq, valid_b_lenseq, valid_additional = valid_X['h_np'], valid_X['b_np'], valid_X['h_seqlen'], valid_X['b_seqlen'],valid_X['additional']
                    valid_h_np, valid_b_np, valid_h_sizes, valid_b_sizes, valid_h_sent_sizes, valid_b_sent_sizes = \
                    valid_X['h_np'], valid_X['b_np'], valid_X[
                        'h_sizes'], valid_X['b_sizes'], valid_X['h_sent_sizes'], valid_X['b_sent_sizes']

                    # print("h_sent_sizes:{} b_sent_sizes:{}".format(valid_h_sent_sizes.shape,valid_b_sent_sizes.shape))

                    batches = self.get_batch_attention(valid_h_np, valid_b_np, valid_h_sizes, valid_b_sizes,
                                                       valid_h_sent_sizes, valid_b_sent_sizes, y_valid)

                    batch_losses = []
                    batch_accuracies = []
                    for (
                    valid_h_batch, valid_b_batch, valid_h_sizes_batch, valid_b_sizes_batch, valid_h_sent_sizes_batch,
                    valid_b_sent_sizes_batch, valid_y_batch) in batches:
                        # print("h_sent_sizes:{} b_sent_sizes:{}".format(valid_h_sent_sizes_batch.shape,
                        # valid_b_sent_sizes_batch.shape))

                        feed_dict_valid = {self._X_head: valid_h_batch, self._X_body: valid_b_batch,
                                           self._X_h_sizes: valid_h_sizes_batch, self._X_b_sizes: valid_b_sizes_batch,
                                           self._X_h_sent_sizes: valid_h_sent_sizes_batch,
                                           self._X_b_sent_sizes: valid_b_sent_sizes_batch, self.y: valid_y_batch}
                        if self.tensorboard_logdir:
                            val_acc_batch, val_loss_batch, summary = sess.run(
                                [self._accuracy, self._loss, self._merged_summary], feed_dict=feed_dict_valid)
                            self._file_writer.add_summary(summary, epoch)
                        else:
                            val_acc_batch, val_loss_batch = sess.run([self._accuracy, self._loss],
                                                                     feed_dict=feed_dict_valid)

                        batch_losses.append(val_loss_batch)
                        batch_accuracies.append(val_acc_batch)

                    val_loss = sum(batch_losses) / len(batch_losses)
                    val_acc = sum(batch_accuracies) / len(batch_accuracies)

                    # feed_dict_valid = {self._X_head: valid_h_np, self._X_body: valid_b_np,self._X_body_length:valid_h_lenseq, self._X_head_length:valid_b_lenseq, self.y: y_valid}

                    # if self.tensorboard_logdir:
                    #     val_acc,val_loss,summary = sess.run([self._accuracy,self._loss,self._merged_summary],feed_dict=feed_dict_valid)
                    #     self._file_writer.add_summary(summary,epoch)
                    # else:
                    #     val_acc,val_loss = sess.run([self._accuracy,self._loss],feed_dict = feed_dict_valid)

                    if self.show_progress:
                        if epoch % self.show_progress == 0:
                            print(
                                "Epoch: {} Current training accuracy: {:.4f} ,Current training loss: {:.6f} Validation Accuracy: {:.4f} Validation Loss{:.6f}".format(
                                    epoch + 1, average_acc, average_loss, val_acc, val_loss))

                    if val_loss < best_acc:
                        best_acc = val_loss
                        checks_without_progress = 0
                        print("accuracy has been improved!")
                        best_parameters = self._get_model_parameters()
                    else:
                        checks_without_progress += 1

                    if checks_without_progress > self.max_checks_without_progress:
                        print("Stopping Early! Loss has not improved in {} epoches".format(
                            self.max_checks_without_progress))
                        break
                else:
                    if self.show_progress:
                        if epoch % self.show_progress == 0:
                            print("Epoch: {} Current training accuracy: {:.4f}".format(epoch + 1, train_acc))

            if best_parameters:
                self._restore_model_parameters(best_parameters)
                return self

    def predict_probabilites(self, X):

        if not self._session:
            raise NotFittedError("This %s instance is not fitted yet" % self.__class__.__name__)

        h_np, b_np, h_sizes, b_sizes, h_sent_sizes, b_sent_sizes = X['h_np'], X['b_np'], X['h_sizes'], X['b_sizes'], X[
            'h_sent_sizes'], X['b_sent_sizes']

        with self._session.as_default() as sess:
            return self._probabilites.eval(
                feed_dict={self._X_head: h_np, self._X_body: b_np, self._X_h_sizes: h_sizes, self._X_b_sizes: b_sizes,
                           self._X_h_sent_sizes: h_sent_sizes, self._X_b_sent_sizes: b_sent_sizes})

import tensorflow as tf
import numpy as np
from deep_models.LSTM import LSTM
from datetime import datetime
from sklearn.exceptions import NotFittedError

he_init = tf.contrib.layers.variance_scaling_initializer()


class HARNN(LSTM):

    def __init__(self, h_max_length=20, b_max_length=200, trainable=False, lstm_layers=2, mlp_layers=1,
                 num_neurons=[128, 128, 32], share_parameters=True, average_pooling=False,
                 optimizer=tf.train.AdamOptimizer,
                 learning_rate=0.001, batch_size=128, activation=tf.nn.relu, initializer=he_init, num_epoch=40,
                 batch_norm_momentum=None, dropout_rate=None,
                 max_check_without_progress=20, show_progress=10, tensorboard_logdir=None, random_state=None,
                 embedding=None, l2_lambda=0.01, word_output_size=64, sent_output_size=64, vocab_size=None):
        LSTM.__init__(self, h_max_length, b_max_length, trainable, lstm_layers, mlp_layers, num_neurons,
                      share_parameters, average_pooling, optimizer, learning_rate, batch_size, activation, initializer,
                      num_epoch, batch_norm_momentum, dropout_rate,
                      max_check_without_progress, show_progress, tensorboard_logdir, random_state, embedding, l2_lambda,
                      vocab_size)

        self.word_output_size = word_output_size
        self.sent_output_size = sent_output_size
        self.vocab_size = vocab_size
        self.embedding_size = 100
        if self.embedding is None and self.vocab_size is None:
            raise Exception("Either embedding or vocab_size must be setted!")

    def gru_cell(self, num_neuron):

        gru = tf.contrib.rnn.GRUCell(num_neuron)
        if self.dropout_rate:
            gru = tf.contrib.rnn.DropoutWrapper(gru, input_keep_prob=self.keep_prob, output_keep_prob=self.keep_prob)
        return gru

    def _bidirectional_rnn(self, inputs, inputs_length, scope=None):

        with tf.variable_scope(scope or 'birnn') as scope:

            if self.lstm_layers == 1:
                rnn_cells = self.gru_cell(self.num_neurons[0])
            elif self.lstm_layers > 1:
                rnn_cells = tf.nn.rnn_cell.MultiRNNCell(
                    [self.gru_cell(i) for i in self.num_neurons[:self.lstm_layers]])

            ((fw_outputs, bw_outputs), (fw_states, bw_states)) = tf.nn.bidirectional_dynamic_rnn(rnn_cells, rnn_cells,
                                                                                                 inputs, inputs_length,
                                                                                                 dtype=tf.float32)
            outputs = tf.concat([fw_outputs, bw_outputs], axis=2)

            if self.lstm_layers > 1:
                fw_states = fw_states[self.lstm_layers - 1]
                bw_states = bw_states[self.lstm_layers - 1]

        return outputs, fw_states, bw_states

    def _attention(self, inputs, output_size, initializer=tf.contrib.layers.xavier_initializer(), activation_fn=tf.tanh,
                   scope=None):

        with tf.variable_scope(scope or 'attention') as scope:
            attetion_context_vector = tf.get_variable(name='attention_context_vector', shape=[output_size],
                                                      initializer=initializer, dtype=tf.float32)
            input_projection = tf.contrib.layers.fully_connected(inputs, output_size, activation_fn=activation_fn,
                                                                 scope=scope)

            vector_attn = tf.reduce_sum(tf.multiply(input_projection, attetion_context_vector), axis=2, keep_dims=True)
            attention_weights = tf.nn.softmax(vector_attn, axis=1)
            weighted_projection = tf.multiply(input_projection, attention_weights)

            outputs = tf.reduce_mean(weighted_projection, axis=1)

            return outputs

    def _hierarichcal_attention(self, inputs, sizes, sent_sizes, scope):

        (doc_size, sent_size, word_size, embed_size) = tf.unstack(tf.shape(inputs))
        print(sizes.get_shape())

        with tf.variable_scope(scope):
            word_level_inputs = tf.reshape(inputs, [doc_size * sent_size, word_size, self.embedding_size])
            word_level_lengths = tf.reshape(sent_sizes, [doc_size * sent_size])

            with tf.variable_scope("word") as scope:
                word_encoder_output, _, _ = self._bidirectional_rnn(word_level_inputs, word_level_lengths, scope)

                with tf.variable_scope("attention") as scope:
                    word_level_output = self._attention(word_encoder_output, self.word_output_size, scope=scope)

                    with tf.variable_scope("dropout") as scope:
                        word_level_output = tf.layers.dropout(word_level_output, rate=self.dropout_rate,
                                                              training=self._training)

            sentence_level_inputs = tf.reshape(word_level_output, (doc_size, sent_size, self.word_output_size))
            with tf.variable_scope("sentence") as scope:
                sentence_encoder_outputs, _, _ = self._bidirectional_rnn(sentence_level_inputs, sizes, scope=scope)

                with tf.variable_scope("attention") as scope:
                    sent_level_output = self._attention(sentence_encoder_outputs, self.sent_output_size, scope=scope)

                    with tf.variable_scope("dropout") as scope:
                        sent_level_output = tf.layers.dropout(sent_level_output, rate=self.dropout_rate)

        return sent_level_output

    def _add_embedding(self, inputs, scope):

        with tf.variable_scope(scope):
            with tf.variable_scope("embedding_lookup") as scope:
                if self.embedding is None and self.vocab_size is not None:
                    embedding = tf.get_variable(
                        name="embedding_matrix",
                        shape=[self.vocab_size, self.embedding_size],
                        initializer=tf.contrib.layers.xavier_initializer(),
                        dtype=tf.float32)
                # else:
                #     embedding = tf.Variable(self.embedding,trainable=self.trainable)
                inputs_embedded = tf.nn.embedding_lookup(embedding, inputs)
                return inputs_embedded

    def _ann(self, head_inputs, body_inputs, h_sizes, b_sizes, h_sent_sizes, b_sent_sizes):

        with tf.variable_scope("embedding_lookup") as scope:
            heads_embeddings = self._add_embedding(head_inputs, scope)
            scope.reuse_variables()
            body_embeddings = self._add_embedding(body_inputs, scope)

        if self.share_parameters:
            with tf.variable_scope("heads_and_bodies") as scope:
                head_sent_level_outputs = self._hierarichcal_attention(heads_embeddings, h_sizes, h_sent_sizes, scope)
                scope.reuse_variables()
                body_sent_level_outputs = self._hierarichcal_attention(body_embeddings, b_sizes, b_sent_sizes, scope)
        else:
            with tf.variable_scope("heads") as scope:
                head_sent_level_outputs = self._hierarichcal_attention(heads_embeddings, h_sizes, h_sent_sizes, scope)
            with tf.variable_scope("bodies") as scope:
                body_sent_level_outputs = self._hierarichcal_attention(body_embeddings, b_sizes, b_sent_sizes, scope)

        output = tf.concat([head_sent_level_outputs, body_sent_level_outputs], axis=1)

        for layer in range(self.mlp_layers):

            if self.dropout_rate:
                output = tf.layers.dropout(output, rate=self.dropout_rate, training=self._training)
            output = tf.layers.dense(output, self.num_neurons[self.lstm_layers + layer], activation=self.activation,
                                     kernel_initializer=self.initializer,
                                     name="hidden{}".format(layer + 1))

            output = self.activation(output, name="hidden{}_out".format(layer + 1))

        return output

    def _construct_graph(self, n_outputs, weights):

        if self.randome_state:
            tf.set_random_seed(self.randome_state)
            np.random.seed(self.randome_state)
        X_heads = tf.placeholder(dtype=tf.int32, shape=[None, None, None], name="X_heads")
        X_bodies = tf.placeholder(dtype=tf.int32, shape=[None, None, None], name="X_bodies")
        X_head_sizes = tf.placeholder(dtype=tf.int32, shape=[None], name="head_sizes")
        X_body_sizes = tf.placeholder(dtype=tf.int32, shape=[None], name="body_sizes")
        X_head_sent_sizes = tf.placeholder(dtype=tf.int32, shape=[None, None], name="head_sent_sizes")
        X_body_sent_sizes = tf.placeholder(dtype=tf.int32, shape=[None, None], name="body_sent_sizes")

        # X_addtional = tf.placeholder(tf.float32,shape=[None,features_length],name="additonal_features")
        y_ = tf.placeholder(tf.int32, shape=[None], name="y")
        weights = tf.constant(weights, tf.float32)

        if self.dropout_rate:
            self._training = tf.placeholder_with_default(False, shape=[], name="training")
            self.keep_prob = tf.cond(self._training, lambda: tf.constant(1 - self.dropout_rate),
                                     lambda: tf.constant(1.0))
        else:
            self._training = None

        pre_output = self._ann(X_heads, X_bodies, X_head_sizes, X_body_sizes, X_head_sent_sizes, X_body_sent_sizes)
        logits = tf.layers.dense(pre_output, n_outputs, kernel_initializer=he_init, name="logits")
        # weighted_logits = tf.multiply(logits,weights)
        probabilities = tf.nn.softmax(logits, name="probabilities")

        xentropy = tf.nn.sparse_softmax_cross_entropy_with_logits(labels=y_, logits=logits)
        loss = tf.reduce_mean(xentropy, name="loss")
        variables = tf.trainable_variables()
        for v in variables:
            print(v.name)
        l2_loss = tf.add_n(
            [tf.nn.l2_loss(v) for v in variables if 'bias' not in v.name and 'Variable' not in v.name]) * self.l2_lambda
        loss += l2_loss

        optimizer = self.optimizer(learning_rate=self.learning_rate)
        update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
        with tf.control_dependencies(update_ops):
            training_op = optimizer.minimize(loss)

        correct = tf.nn.in_top_k(logits, y_, 1)
        accuracy = tf.reduce_mean(tf.cast(correct, tf.float32), name="accuracy")

        init = tf.global_variables_initializer()
        saver = tf.train.Saver()

        if self.tensorboard_logdir:
            now = datetime.utcnow().strftime('%Y%m%d-%H%M%S')
            tb_logdir = self.tensorboard_logdir + "/run{}".format(now)
            cost_summary = tf.summary.scalar("validation_loss", loss)
            acc_summary = tf.summary.scalar("validation_accuracy", accuracy)
            merged_summary = tf.summary.merge_all()
            file_writer = tf.summary.FileWriter(tb_logdir, tf.get_default_graph())

            self._merged_summary = merged_summary
            self._file_writer = file_writer

        self._X_head, self._X_body, self._X_h_sizes, self._X_b_sizes, self._X_h_sent_sizes, self._X_b_sent_sizes, self.y = X_heads, X_bodies, X_head_sizes, X_body_sizes, X_head_sent_sizes, X_body_sent_sizes, y_
        self._logits = logits
        self._probabilites = probabilities
        self._loss = loss
        self._training_op = training_op
        self._accuracy = accuracy
        self._init, self._saver = init, saver
        # self._X_additional = X_addtional

    def get_batch_attention(self, h_np, b_np, h_sizes, b_sizes, h_sent_sizes, b_sent_sizes, y):

        num_batches = h_np.shape[0] // self.batch_size
        batches = []

        for i in range(num_batches):
            if (i + 1) * self.batch_size > h_np.shape[0]:
                batch_h_np = h_np[i * self.batch_size:h_np.shape[0] - 1, :]
                batch_b_np = b_np[i * self.batch_size:h_np.shape[0] - 1, :]
                batch_h_sizes = h_sizes[i * self.batch_size:h_np.shape[0] - 1]
                batch_b_sizes = b_sizes[i * self.batch_size:h_np.shape[0] - 1]
                batch_h_sent_sizes = h_sent_sizes[i * self.batch_size:h_np.shape[0] - 1, :]
                batch_b_sent_sizes = b_sent_sizes[i * self.batch_size:h_np.shape[0] - 1, :]
                batch_y = y[i * self.batch_size:h_np.shape[0] - 1]
                # batch_additional = additional_features[i*self.batch_size:h_np.shape[0]-1]
            else:
                batch_h_np = h_np[i * self.batch_size:(i + 1) * self.batch_size, :]
                batch_b_np = b_np[i * self.batch_size:(i + 1) * self.batch_size, :]
                batch_h_sizes = h_sizes[i * self.batch_size:(i + 1) * self.batch_size]
                batch_b_sizes = b_sizes[i * self.batch_size:(i + 1) * self.batch_size]
                batch_h_sent_sizes = h_sent_sizes[i * self.batch_size:(i + 1) * self.batch_size, :]
                batch_b_sent_sizes = b_sent_sizes[i * self.batch_size:(i + 1) * self.batch_size, :]
                batch_y = y[i * self.batch_size:(i + 1) * self.batch_size]
                # batch_additional = additional_features[i*self.batch_size:(i+1)*self.batch_size]
            batches.append(
                (batch_h_np, batch_b_np, batch_h_sizes, batch_b_sizes, batch_h_sent_sizes, batch_b_sent_sizes, batch_y))

        return batches

    def sort_base_bodies(self, h_v, b_v, y):

        assert len(h_v) == len(b_v) == len(y)
        b_lengths = [len(b_vector) for b_vector in b_v]

        sorted_h = []
        sorted_b = []
        sorted_y = []
        for length in range(min(b_lengths), max(b_lengths)):
            for idx, vectors in enumerate(b_v):
                if len(vectors) == length:
                    sorted_h.append(h_v[idx])
                    sorted_b.append(b_v[idx])
                    sorted_y.append(y[idx])

        return sorted_h, sorted_b, sorted_y

    def get_batch(self, h_v, b_v, y=None):

        assert len(h_v) == len(b_v)

        for batch_i in range(0, len(b_v) // self.batch_size):
            start_i = batch_i * self.batch_size
            h_batch = h_v[start_i:start_i + self.batch_size]
            b_batch = b_v[start_i:start_i + self.batch_size]
            if y is not None:
                y_batch = y[start_i:start_i + self.batch_size]

            h_lengths_batch = [len(h) for h in h_batch]
            b_lengths_batch = [len(b) for b in b_batch]

            def padding(vectors, lengths):

                max_length = max(lengths)
                np_batch = np.zeros(shape=(len(vectors), max_length, 4096))
                for i, sents_vector in enumerate(vectors):
                    for j, sent in enumerate(sents_vector):
                        np_batch[i, j, :] = sent
                return np_batch

            np_h_batch = padding(h_batch, h_lengths_batch)
            np_b_batch = padding(b_batch, b_lengths_batch)
            h_lengths_batch = np.asarray(h_lengths_batch, np.int32)
            b_lengths_batch = np.asarray(b_lengths_batch, np.int32)

            if y is not None:
                yield np_h_batch, np_b_batch, h_lengths_batch, b_lengths_batch, y_batch
            else:
                yield np_h_batch, np_b_batch, h_lengths_batch, b_lengths_batch

    def fit(self, X, y, valid_X=None, y_valid=None):
        self.close_session()
        y = np.array(y)

        if len(y.shape) == 2:
            y = np.argmax(y, axis=1)

        self._classes, counts = np.unique(y, return_counts=True)
        weights = [count / sum(counts) for count in counts]
        weights = np.asarray(weights, np.float32)
        n_outputs = len(self._classes)

        self._graph = tf.Graph()
        h_np, b_np, h_sizes, b_sizes, h_sent_sizes, b_sent_sizes = X['h_np'], X['b_np'], X['h_sizes'], X['b_sizes'], X[
            'h_sent_sizes'], X['b_sent_sizes']
        #        features_length = addtional_features.shape[1]

        # print("head shape{}".format(h_sizes.shape))

        with self._graph.as_default():
            self._construct_graph(n_outputs, weights)

        checks_without_progress = 0
        best_acc = np.inf
        best_parameters = None
        gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.333)
        self._session = tf.Session(graph=self._graph, config=tf.ConfigProto(gpu_options=gpu_options))

        with self._session.as_default() as sess:
            self._init.run()
            num_instances = h_np.shape[0]
            for epoch in range(self.num_epoch):
                losses = []
                accs = []
                rnd_idx = np.random.permutation(num_instances)
                for rnd_indices in np.array_split(rnd_idx, num_instances // self.batch_size):

                    X_head_batch, X_body_batch, X_h_sizes_batch, X_b_sizes_batch, X_h_sent_sizes_batch, X_b_sent_sizes_batch, y_batch = \
                    h_np[rnd_indices], b_np[rnd_indices], \
                    h_sizes[rnd_indices], b_sizes[rnd_indices], \
                    h_sent_sizes[rnd_indices], b_sent_sizes[rnd_indices], \
                    y[rnd_indices]
                    # print("batch head shape{}".format(X_h_sizes_batch.shape))
                    y_batch = np.asarray(y_batch)
                    feed_dict = {self._X_head: X_head_batch, self._X_body: X_body_batch,
                                 self._X_h_sizes: X_h_sizes_batch, self._X_b_sizes: X_b_sizes_batch,
                                 self._X_h_sent_sizes: X_h_sent_sizes_batch, self._X_b_sent_sizes: X_b_sent_sizes_batch,
                                 self.y: y_batch}
                    if self._training is not None:
                        feed_dict[self._training] = True

                    train_acc, _, loss = sess.run([self._accuracy, self._training_op, self._loss], feed_dict=feed_dict)
                    losses.append(loss)
                    accs.append(train_acc)
                average_loss = sum(losses) / len(losses)
                average_acc = sum(accs) / len(accs)

                if valid_X is not None and y_valid is not None:

                    print("validation phase")
                    # valid_h_np, valid_b_np, valid_h_lenseq, valid_b_lenseq, valid_additional = valid_X['h_np'], valid_X['b_np'], valid_X['h_seqlen'], valid_X['b_seqlen'],valid_X['additional']
                    valid_h_np, valid_b_np, valid_h_sizes, valid_b_sizes, valid_h_sent_sizes, valid_b_sent_sizes = \
                    valid_X['h_np'], valid_X['b_np'], valid_X[
                        'h_sizes'], valid_X['b_sizes'], valid_X['h_sent_sizes'], valid_X['b_sent_sizes']

                    # print("h_sent_sizes:{} b_sent_sizes:{}".format(valid_h_sent_sizes.shape,valid_b_sent_sizes.shape))

                    batches = self.get_batch_attention(valid_h_np, valid_b_np, valid_h_sizes, valid_b_sizes,
                                                       valid_h_sent_sizes, valid_b_sent_sizes, y_valid)

                    batch_losses = []
                    batch_accuracies = []
                    for (
                    valid_h_batch, valid_b_batch, valid_h_sizes_batch, valid_b_sizes_batch, valid_h_sent_sizes_batch,
                    valid_b_sent_sizes_batch, valid_y_batch) in batches:
                        # print("h_sent_sizes:{} b_sent_sizes:{}".format(valid_h_sent_sizes_batch.shape,
                        # valid_b_sent_sizes_batch.shape))

                        feed_dict_valid = {self._X_head: valid_h_batch, self._X_body: valid_b_batch,
                                           self._X_h_sizes: valid_h_sizes_batch, self._X_b_sizes: valid_b_sizes_batch,
                                           self._X_h_sent_sizes: valid_h_sent_sizes_batch,
                                           self._X_b_sent_sizes: valid_b_sent_sizes_batch, self.y: valid_y_batch}
                        if self.tensorboard_logdir:
                            val_acc_batch, val_loss_batch, summary = sess.run(
                                [self._accuracy, self._loss, self._merged_summary], feed_dict=feed_dict_valid)
                            self._file_writer.add_summary(summary, epoch)
                        else:
                            val_acc_batch, val_loss_batch = sess.run([self._accuracy, self._loss],
                                                                     feed_dict=feed_dict_valid)

                        batch_losses.append(val_loss_batch)
                        batch_accuracies.append(val_acc_batch)

                    val_loss = sum(batch_losses) / len(batch_losses)
                    val_acc = sum(batch_accuracies) / len(batch_accuracies)

                    # feed_dict_valid = {self._X_head: valid_h_np, self._X_body: valid_b_np,self._X_body_length:valid_h_lenseq, self._X_head_length:valid_b_lenseq, self.y: y_valid}

                    # if self.tensorboard_logdir:
                    #     val_acc,val_loss,summary = sess.run([self._accuracy,self._loss,self._merged_summary],feed_dict=feed_dict_valid)
                    #     self._file_writer.add_summary(summary,epoch)
                    # else:
                    #     val_acc,val_loss = sess.run([self._accuracy,self._loss],feed_dict = feed_dict_valid)

                    if self.show_progress:
                        if epoch % self.show_progress == 0:
                            print(
                                "Epoch: {} Current training accuracy: {:.4f} ,Current training loss: {:.6f} Validation Accuracy: {:.4f} Validation Loss{:.6f}".format(
                                    epoch + 1, average_acc, average_loss, val_acc, val_loss))

                    if val_loss < best_acc:
                        best_acc = val_loss
                        checks_without_progress = 0
                        print("accuracy has been improved!")
                        best_parameters = self._get_model_parameters()
                    else:
                        checks_without_progress += 1

                    if checks_without_progress > self.max_checks_without_progress:
                        print("Stopping Early! Loss has not improved in {} epoches".format(
                            self.max_checks_without_progress))
                        break
                else:
                    if self.show_progress:
                        if epoch % self.show_progress == 0:
                            print("Epoch: {} Current training accuracy: {:.4f}".format(epoch + 1, train_acc))

            if best_parameters:
                self._restore_model_parameters(best_parameters)
                return self

    def predict_probabilites(self, X):

        if not self._session:
            raise NotFittedError("This %s instance is not fitted yet" % self.__class__.__name__)

        h_np, b_np, h_sizes, b_sizes, h_sent_sizes, b_sent_sizes = X['h_np'], X['b_np'], X['h_sizes'], X['b_sizes'], X[
            'h_sent_sizes'], X['b_sent_sizes']

        with self._session.as_default() as sess:
            return self._probabilites.eval(
                feed_dict={self._X_head: h_np, self._X_body: b_np, self._X_h_sizes: h_sizes, self._X_b_sizes: b_sizes,
                           self._X_h_sent_sizes: h_sent_sizes, self._X_b_sent_sizes: b_sent_sizes})

# else: #     embedding = tf.Variable(self.embedding,trainable=self.trainable) # X_addtional = tf.placeholder(tf.float32,shape=[None,features_length],name="additonal_features") # weighted_logits = tf.multiply(logits,weights) # self._X_additional = X_addtional # batch_additional = additional_features[i*self.batch_size:h_np.shape[0]-1] # batch_additional = additional_features[i*self.batch_size:(i+1)*self.batch_size] #        features_length = addtional_features.shape[1] # print("head shape{}".format(h_sizes.shape)) # print("batch head shape{}".format(X_h_sizes_batch.shape)) # valid_h_np, valid_b_np, valid_h_lenseq, valid_b_lenseq, valid_additional = valid_X['h_np'], valid_X['b_np'], valid_X['h_seqlen'], valid_X['b_seqlen'],valid_X['additional'] # print("h_sent_sizes:{} b_sent_sizes:{}".format(valid_h_sent_sizes.shape,valid_b_sent_sizes.shape)) # print("h_sent_sizes:{} b_sent_sizes:{}".format(valid_h_sent_sizes_batch.shape, # valid_b_sent_sizes_batch.shape)) # feed_dict_valid = {self._X_head: valid_h_np, self._X_body: valid_b_np,self._X_body_length:valid_h_lenseq, self._X_head_length:valid_b_lenseq, self.y: y_valid} # if self.tensorboard_logdir: #     val_acc,val_loss,summary = sess.run([self._accuracy,self._loss,self._merged_summary],feed_dict=feed_dict_valid) #     self._file_writer.add_summary(summary,epoch) # else: #     val_acc,val_loss = sess.run([self._accuracy,self._loss],feed_dict = feed_dict_valid)