""" Code for the MAML algorithm and network architecture. """

import numpy as np

import sklearn

import tensorflow as tf

import os, time, shutil, collections

from tensorflow.contrib import rnn

import tensorflow.contrib.layers as layers

from tensorflow.contrib.rnn import RNNCell

from tensorflow.python.platform import flags

FLAGS = flags.FLAGS

PADDING_ID = 1016

WORDS_NUM = 1017

MASK_ARRAY = [[1.]] * PADDING_ID + [[0.]] + [[1.]] * (WORDS_NUM - PADDING_ID - 1)

SUMMARY_INTERVAL = 100

SAVE_INTERVAL = 1000

PRINT_INTERVAL = 100

TEST_PRINT_INTERVAL = PRINT_INTERVAL*5

class BaseModel(object):

    """

    Base Model for basic networks with sequential data, i.e., RNN, CNN.

    """

    def __init__(self):

        self.regularizers = []

        self.regularization = 0.01

        self.isReg = True

    def evaluate(self, data, labels, sess=None, prefix="metatest_"):

        """

        Runs one evaluation against the full epoch of data.

        Return the precision and the number of correct predictions.

        Batch evaluation saves memory and enables this to run on smaller GPUs.

        sess: the session in which the model has been trained.

        op: the Tensor that returns the number of correct predictions.

        """

        t_process, t_wall = time.process_time(), time.time()

        predictions, loss = self.predict(data, labels, sess)

        fpr, tpr, _ = sklearn.metrics.roc_curve(labels, predictions)

        auc = 100 * sklearn.metrics.auc(fpr, tpr)

        ncorrects = sum(predictions == labels)

        accuracy = 100 * sklearn.metrics.accuracy_score(labels, predictions)

        string = 'auc: {:.2f}, accuracy: {:.2f} ({:d} / {:d}), loss: {:.2e}'.format(auc, accuracy, ncorrects, len(labels), loss)

        if sess is None:

            string += '\ntime: {:.0f}s (wall {:.0f}s)'.format(time.process_time()-t_process, time.time()-t_wall)

        # return string, auc, loss, predictions

        return string, auc, accuracy, loss, predictions

    def fit(self, X_tr, y_tr, X_vl, y_vl):

        t_process, t_wall = time.process_time(), time.time()

        sess = tf.Session(graph=self.graph)

        shutil.rmtree(self._get_path('summaries'), ignore_errors=True)

        writer = tf.summary.FileWriter(self._get_path('summaries'), self.graph)

        shutil.rmtree(self._get_path('checkpoints'), ignore_errors=True)

        os.makedirs(self._get_path('checkpoints'))

        path = os.path.join(self._get_path('checkpoints'), 'model')

        sess.run(self.op_init)

        # Training.

        count = 0

        bad_counter = 0

        accuracies = []

        aucs = []

        losses = []

        indices = collections.deque()

        num_steps = int(self.num_epochs * X_tr.shape[0] / self.batch_size)

        estop = False  # early stop

        if type(X_vl) is not np.ndarray:

            X_vl = X_vl.toarray()

        for step in range(1, num_steps+1):

            # Be sure to have used all the samples before using one a second time.

            if len(indices) < self.batch_size:

                indices.extend(np.random.permutation(X_tr.shape[0]))

            idx = [indices.popleft() for i in range(self.batch_size)]

            count += len(idx)

            batch_data, batch_labels = X_tr[idx, :, :], y_tr[idx]

            if type(batch_data) is not np.ndarray:

                batch_data = batch_data.toarray()  # convert sparse matrices

            feed_dict = {self.ph_data: batch_data, self.ph_labels: batch_labels, self.ph_dropout: self.dropout, self.ph_training: True}

            learning_rate, loss_average = sess.run([self.op_train, self.op_loss_average], feed_dict)

            # Periodical evaluation of the model.

            if step % self.eval_frequency == 0 or step == num_steps:

                print ('Seen samples: %d' % count)

                epoch = step * self.batch_size / X_tr.shape[0]

                print('step {} / {} (epoch {:.2f} / {}):'.format(step, num_steps, epoch, self.num_epochs))

                print('  learning_rate = {:.2e}, loss_average = {:.2e}'.format(learning_rate, loss_average))

                string, auc, accuracy, loss, predictions = self.evaluate(X_vl, y_vl, sess)

                aucs.append(auc)

                accuracies.append(accuracy)

                losses.append(loss)

                print('  validation {}'.format(string))

                print('  time: {:.0f}s (wall {:.0f}s)'.format(time.process_time()-t_process, time.time()-t_wall))

                # Summaries for TensorBoard.

                summary = tf.Summary()

                summary.ParseFromString(sess.run(self.op_summary, feed_dict))

                summary.value.add(tag='validataion/auc', simple_value=auc)

                summary.value.add(tag='validation/loss', simple_value=loss)

                writer.add_summary(summary, step)

                # Save model parameters (for evaluation).

                self.op_saver.save(sess, path, global_step=step)

                if len(aucs) > (self.patience+5) and auc > np.array(aucs).max():

                    bad_counter = 0

                if len(aucs) > (self.patience+5) and auc <= np.array(aucs)[:-self.patience].max():

                    bad_counter += 1

                    if bad_counter > self.patience:

                        print('Early Stop!')

                        estop = True

                        break

            if estop:

                break

        print('validation accuracy: peak = {:.2f}, mean = {:.2f}'.format(max(accuracies), np.mean(accuracies[-10:])))

        print('validation auc: peak = {:.2f}, mean = {:.2f}'.format(max(aucs), np.mean(aucs[-10:])))

        # store weights value for fine-tune

        if self.is_finetune is not True:

            feed_dict = {}

            for k in self.op_weights:

                self.weights_for_init[k] = sess.run([self.op_weights[k]], feed_dict)[0]

                self.weights_for_finetune[k] = sess.run([self.op_weights[k]], feed_dict)[0]

        writer.close()

        sess.close()

        t_step = (time.time() - t_wall) / num_steps

        return  sess, aucs, accuracies

    def loss(self, logits):

        # Define loss and optimizer

        with tf.name_scope('cross_entropy'):

            labels = tf.to_int64(self.ph_labels)

            cross_entropy = tf.nn.sparse_softmax_cross_entropy_with_logits(logits=logits, labels=labels)

            cross_entropy = tf.reduce_mean(cross_entropy)

        if self.is_finetune and self.freeze_opt == 'mlp':

            loss = cross_entropy

            # Summaries for TensorBoard.

            tf.summary.scalar('loss/cross_entropy', cross_entropy)

            tf.summary.scalar('loss/total', loss)

            with tf.name_scope('averages'):

                averages = tf.train.ExponentialMovingAverage(0.9)

                op_averages = averages.apply([cross_entropy, loss])

                tf.summary.scalar('loss/avg/cross_entropy', averages.average(cross_entropy))

                tf.summary.scalar('loss/avg/total', averages.average(loss))

                with tf.control_dependencies([op_averages]):

                    loss_average = tf.identity(averages.average(loss), name='control')

        else:

            with tf.name_scope('regularization'):

                regularization = self.regularization

                regularization *= tf.add_n(self.regularizers)

            loss = cross_entropy + regularization

            # Summaries for TensorBoard.

            tf.summary.scalar('loss/cross_entropy', cross_entropy)

            tf.summary.scalar('loss/regularization', regularization)

            tf.summary.scalar('loss/total', loss)

            with tf.name_scope('averages'):

                averages = tf.train.ExponentialMovingAverage(0.9)

                op_averages = averages.apply([cross_entropy, regularization, loss])

                tf.summary.scalar('loss/avg/cross_entropy', averages.average(cross_entropy))

                tf.summary.scalar('loss/avg/regularization', averages.average(regularization))

                tf.summary.scalar('loss/avg/total', averages.average(loss))

                with tf.control_dependencies([op_averages]):

                    loss_average = tf.identity(averages.average(loss), name='control')

        return loss, loss_average

    def predict(self, data, labels=None, sess=None):

        loss = 0

        size = data.shape[0]

        predictions = np.empty(size)

        sess = self._get_session(sess)

        for begin in range(0, size, self.batch_size):

            end = begin + self.batch_size

            end = min([end, size])

            batch_data = np.zeros((self.batch_size, data.shape[1], data.shape[2]))

            tmp_data = data[begin:end, :, :]

            if type(tmp_data) is not np.ndarray:

                tmp_data = tmp_data.toarray()  # convert sparse matrices

            batch_data[:end-begin] = tmp_data

            feed_dict = {self.ph_data: batch_data, self.ph_dropout: 1, self.ph_training: False}

            # Compute loss if labels are given.

            if labels is not None:

                batch_labels = np.zeros(self.batch_size)

                batch_labels[:end-begin] = labels[begin:end]

                feed_dict[self.ph_labels] = batch_labels

                batch_pred, batch_loss = sess.run([self.op_prediction, self.op_loss], feed_dict)

                loss += batch_loss

            else:

                batch_pred = sess.run(self.op_prediction, feed_dict)

            predictions[begin:end] = batch_pred[:end-begin]

        if labels is not None:

            return predictions, loss * self.batch_size / size

        else:

            return predictions

    def training(self, loss, learning_rate, decay_steps, decay_rate=0.95, momentum=0.9):

        """Adds to the loss model the Ops required to generate and apply gradients."""

        with tf.name_scope('training'):

            # Learning rate.

            global_step = tf.Variable(0, name='global_step', trainable=False)

            if decay_rate != 1:

                learning_rate = tf.train.exponential_decay(

                        learning_rate, global_step, decay_steps, decay_rate, staircase=True)

            tf.summary.scalar('learning_rate', learning_rate)

            # Optimizer.

            if momentum == 0:

                optimizer = tf.train.GradientDescentOptimizer(learning_rate)

            else:

                optimizer = tf.train.MomentumOptimizer(learning_rate, momentum)

            grads = optimizer.compute_gradients(loss)

            op_gradients = optimizer.apply_gradients(grads, global_step=global_step)

            # Histograms.

            for grad, var in grads:

                if grad is None:

                    print('warning: {} has no gradient'.format(var.op.name))

                else:

                    tf.summary.histogram(var.op.name + '/gradients', grad)

            # The op return the learning rate.

            with tf.control_dependencies([op_gradients]):

                op_train = tf.identity(learning_rate, name='control')

        return op_train

     # Helper methods.

    def _get_path(self, folder):

        path = '../../models/'

        return os.path.join(path, folder, self.dir_name)

    def _get_session(self, sess=None):

        """Restore parameters if no session given."""

        if sess is None:

            sess = tf.Session(graph=self.graph)

            filename = tf.train.latest_checkpoint(self._get_path('checkpoints'))

            self.op_saver.restore(sess, filename)

        return sess

    def _get_prediction(self, logits):

        """Return the predicted classes."""

        with tf.name_scope('prediction'):

            prediction = tf.argmax(logits, axis=1)

            return prediction

     # Helper methods.

    def _get_path(self, folder):

        path = '../../models/'

        return os.path.join(path, folder, self.dir_name)

    def _get_session(self, sess=None):

        """Restore parameters if no session given."""

        if sess is None:

            sess = tf.Session(graph=self.graph)

            filename = tf.train.latest_checkpoint(self._get_path('checkpoints'))

            self.op_saver.restore(sess, filename)

        return sess

    def weight_variable(self, shape, name='weights'):

        initial = tf.truncated_normal_initializer(0, 0.1)

        var = tf.get_variable(name, shape, tf.float32, initializer=initial)

        if self.isReg:

            self.regularizers.append(tf.nn.l2_loss(var))

        tf.summary.histogram(var.op.name, var)

        return var

    def bias_variable(self, shape, name='bias'):

        initial = tf.constant_initializer(0.1)

        var = tf.get_variable(name, shape, tf.float32, initializer=initial)

        if self.isReg:

            self.regularizers.append(tf.nn.l2_loss(var))

        tf.summary.histogram(var.op.name, var)

        return var

    def build_fc_weights(self, dim_in, weights):

        for i, dim in enumerate(self.dim_hidden):

            dim_out = dim

            weights["fc_W"+str(i)] = self.weight_variable([int(dim_in), dim_out], name="fc_W"+str(i))

            weights["fc_b"+str(i)] = self.bias_variable([dim_out], name="fc_b"+str(i))

            dim_in = dim_out

        return weights

    def fc(self, x, W, b, relu=True):

        """Fully connected layer with Mout features."""

        x = tf.matmul(x, W) + b

        return tf.nn.relu(x) if relu else x

    def normalize(self, inputs, epsilon = 1e-8, scope="ln", reuse=None):

        '''Applies layer normalization.

        Args:

          inputs: A tensor with 2 or more dimensions, where the first dimension has

            `batch_size`.

          epsilon: A floating number. A very small number for preventing ZeroDivision Error.

          scope: Optional scope for `variable_scope`.

          reuse: Boolean, whether to reuse the weights of a previous layer

            by the same name.

        Returns:

          A tensor with the same shape and data dtype as `inputs`.

        '''

        with tf.variable_scope(scope, reuse=reuse):

            inputs_shape = inputs.get_shape()

            params_shape = inputs_shape[-1:]

            mean, variance = tf.nn.moments(inputs, [-1], keep_dims=True)

            beta= tf.Variable(tf.zeros(params_shape))

            gamma = tf.Variable(tf.ones(params_shape))

            normalized = (inputs - mean) / ( (variance + epsilon) ** (.5) )

            outputs = gamma * normalized + beta

        return outputs

class RNN(BaseModel):

    """

    Build a vanilla recurrent neural network.

    """

    def __init__(self, data_loader, weights_for_finetune, init_std=0.05, freeze_opt=None, is_finetune=False):

        super().__init__()

        self.is_finetune = is_finetune

        self.freeze_opt = freeze_opt

        print ("freeze_opt: ", self.freeze_opt)

        if self.is_finetune:

            self.finetune_weights = weights_for_finetune

            self.learning_rate = 0.00001

            self.batch_size = 128

            self.num_epochs = 30

        else:

            self.learning_rate = 0.5

            self.batch_size = 128

            self.num_epochs = 200

        # training parameters

        self.dir_name =  "rnn"

        self.dropout = 1

        self.decay_rate = 0.9

        self.decay_steps = 10000 / self.batch_size

        self.momentum = 0.95

        self.patience = 5

        self.eval_frequency = self.num_epochs

        # Network Parameters

        self.init_std = init_std

        self.n_hidden = 256 # hidden dimensions of embedding

        self.n_hidden_1 = 128

        self.n_hidden_2 = 128

        self.n_words = data_loader.n_words

        self.num_input = data_loader.dim_input

        self.n_classes = FLAGS.n_classes

        self.timesteps = data_loader.timesteps

        self.code_size = data_loader.code_size

        self.dim_hidden = [self.n_hidden_1, self.n_hidden_2, FLAGS.n_classes]

        self.weights_for_init = dict() # to store the value of learned params

        self.weights_for_finetune = dict()

        self.build_model()

        print('method', self.dir_name, 'data shape:', self.num_input, 'batch size:', self.batch_size, 'learning rate:', self.learning_rate, \

              'momentum:', self.momentum, 'patience:', self.patience)

    # Methods to construct the computational graph

    def build_model(self):

        """Build the computational graph with memory network of the model."""

        self.graph = tf.Graph()

        with self.graph.as_default():

            # Inputs.

            with tf.name_scope('inputs'):

                # tf Graph input

                self.ph_data = tf.placeholder(tf.int32, (self.batch_size, self.timesteps, self.code_size), 'data')

                self.ph_labels = tf.placeholder(tf.int32, (self.batch_size), 'labels')

                self.ph_dropout = tf.placeholder(tf.float32, (), 'dropout')

                self.ph_training = tf.placeholder(tf.bool, name='trainingFlag')

            # Construct model

            op_logits = self._inference(self.ph_data, self.ph_dropout, self.ph_training)

            self.op_loss, self.op_loss_average = self.loss(op_logits)

            self.op_train = self.training(self.op_loss, self.learning_rate,

                    self.decay_steps, self.decay_rate, self.momentum)

            self.op_prediction = self._get_prediction(op_logits)

            # Initialize variables, i.e. weights and biases.

            self.op_init = tf.global_variables_initializer()

            if self.is_finetune is not True:

                self.op_weights = self.get_op_variables()

            else:

                print (tf.trainable_variables())

            # Summaries for TensorBoard and Save for model parameters.

            self.op_summary = tf.summary.merge_all()

            self.op_saver = tf.train.Saver(max_to_keep=5)

        self.graph.finalize()

    def get_op_variables(self):

        op_weights = dict()

        op_var = tf.trainable_variables()

        # embedding

        op_weights["emb_W"] = [v for v in op_var if "emb_W" in v.name][0]

        # lstm

        op_weights["lstm_W_xh"] = [v for v in op_var if "lstm_W_xh" in v.name][0]

        op_weights["lstm_W_hh"] = [v for v in op_var if "lstm_W_hh" in v.name][0]

        op_weights["lstm_b"] = [v for v in op_var if "lstm_b" in v.name][0]

        # fully connected

        for i, dim in enumerate(self.dim_hidden):

            op_weights["fc_W"+str(i)] = [v for v in op_var if "fc_W"+str(i) in v.name ][0]

            op_weights["fc_b"+str(i)] = [v for v in op_var if "fc_b"+str(i) in v.name][0]

        print ('show variable')

        print(op_var)

        return op_weights

    def build_emb_weights(self, weights):

        weights["emb_W"] = tf.Variable(tf.random_normal([self.n_words, self.n_hidden], stddev=self.init_std), name="emb_W")

        weights["emb_mask_W"] = tf.get_variable("mask_padding", initializer=MASK_ARRAY, dtype="float32", trainable=False)

        return weights

    def embedding(self, x, Wemb, Wemb_mask):

        _x = tf.nn.embedding_lookup(Wemb, x) # recs size is (batch_size, timesteps, n_words)

        _x_mask = tf.nn.embedding_lookup(Wemb_mask, x)

        emb_vecs = tf.multiply(_x, _x_mask) # broadcast

        emb_vecs = tf.reduce_sum(emb_vecs, 2)

        return emb_vecs

    def lstm_identity_initializer(self, scale):

        def _initializer(shape, dtype=tf.float32, partition_info=None):

            """Ugly cause LSTM params calculated in one matrix multiply"""

            size = shape[0]

            # gate (j) is identity

            t = np.zeros(shape)

            t[:, size:size * 2] = np.identity(size) * scale

            t[:, :size] = self.orthogonal([size, size])

            t[:, size * 2:size * 3] = self.orthogonal([size, size])

            t[:, size * 3:] = self.orthogonal([size, size])

            return tf.constant(t, dtype=dtype)

        return _initializer

    def orthogonal_initializer(self):

        def _initializer(shape, dtype=tf.float32, partition_info=None):

            return tf.constant(self.orthogonal(shape), dtype)

        return _initializer

    def orthogonal(self, shape):

        flat_shape = (shape[0], np.prod(shape[1:]))

        a = np.random.normal(0.0, 1.0, flat_shape)

        u, _, v = np.linalg.svd(a, full_matrices=False)

        q = u if u.shape == flat_shape else v

        return q.reshape(shape)

    def build_lstm_weights(self, weights):

        #

        # # Keep W_xh and W_hh separate here as well to reuse initialization methods

        # with tf.variable_scope(scope or type(self).__name__):

        weights["lstm_W_xh"] = tf.get_variable('lstm_W_xh', [self.n_hidden, 4 * self.n_hidden],

                               initializer=self.orthogonal_initializer())

        weights["lstm_W_hh"] = tf.get_variable('lstm_W_hh', [self.n_hidden, 4 * self.n_hidden],

                               initializer=self.lstm_identity_initializer(0.95),)

        weights["lstm_b"] = tf.get_variable('lstm_b', [4 * self.n_hidden])

        return weights

    # Create model

    def _inference(self, x, dropout, is_training=True):

        with tf.variable_scope('pretrain_model', reuse=None) as training_scope:

            if self.freeze_opt == None:

                weights = {}

                weights = self.build_emb_weights(weights)

                weights = self.build_lstm_weights(weights)

                weights = self.build_fc_weights(self.n_hidden, weights)

                # embedding

                with tf.variable_scope("embedding"):

                    xemb = self.embedding(x, weights["emb_W"], weights["emb_mask_W"])

                # recurrent neural networks

                with tf.variable_scope("rnn"):

                    lstm_cell = LSTMCell(self.n_hidden, weights["lstm_W_xh"], weights["lstm_W_hh"], weights["lstm_b"])

                    # lstm_cell = LSTMCell(self.n_hidden)

                    xemb = tf.unstack(xemb, self.timesteps, 1)

                    #c, h

                    W_state_c = tf.random_normal([self.batch_size, self.n_hidden], stddev=0.1)

                    W_state_h = tf.random_normal([self.batch_size, self.n_hidden], stddev=0.1)

                    outputs, state = tf.nn.static_rnn(lstm_cell, xemb, initial_state=(W_state_c, W_state_h), dtype=tf.float32)

                    _, hout = state

                with tf.variable_scope("dropout"):

                    h_ = layers.dropout(hout, keep_prob=dropout)

                for i, dim in enumerate(self.dim_hidden[:-1]):

                    h_ = self.fc(h_, weights["fc_W"+str(i)], weights["fc_b"+str(i)])

                    h_ = tf.nn.dropout(h_, dropout)

                # Logits linear layer, i.e. softmax without normalization.

                N, Min = h_.get_shape()

                i = len(self.dim_hidden)-1

                logits = self.fc(h_, weights["fc_W"+str(i)], weights["fc_b"+str(i)], relu=False)

            else:

                with tf.variable_scope("embedding"):

                    Wemb = self.finetune_weights["emb_W"]

                    Wemb_mask = tf.get_variable("mask_padding", initializer=MASK_ARRAY, dtype="float32", trainable=False)

                    xemb = self.embedding(x, Wemb, Wemb_mask)

                # convolutional network

                with tf.variable_scope("rnn"):

                    lstm_cell = LSTMCell(self.n_hidden, self.finetune_weights["lstm_W_xh"], self.finetune_weights["lstm_W_hh"], self.finetune_weights["lstm_b"])

                    xemb = tf.unstack(xemb, self.timesteps, 1)

                    W_state_c = tf.random_normal([self.batch_size, self.n_hidden], stddev=0.1)

                    W_state_h = tf.random_normal([self.batch_size, self.n_hidden], stddev=0.1)

                    outputs, state = tf.nn.static_rnn(lstm_cell, xemb, initial_state=(W_state_c, W_state_h), dtype=tf.float32)

                    _, hout = state

                with tf.variable_scope("dropout"):

                    h_ = layers.dropout(hout, keep_prob=dropout)

                for i, dim in enumerate(self.dim_hidden[:-1]):

                    Wfc = self.finetune_weights["fc_W"+str(i)]

                    bfc = self.finetune_weights["fc_b"+str(i)]

                    h_ = self.fc(h_, Wfc, bfc)

                    h_ = tf.nn.dropout(h_, dropout)

                # finetune the last layer

                i = len(self.dim_hidden)-1

                weights = {}

                dim_in = self.n_hidden_2

                weights["fc_W"+str(i)] = self.weight_variable([int(dim_in), FLAGS.n_classes], name="fc_W"+str(i))

                weights["fc_b"+str(i)] = self.bias_variable([FLAGS.n_classes], name="fc_b"+str(i))

                # Logits linear layer, i.e. softmax without normalization.

                N, Min = h_.get_shape()

                i = len(self.dim_hidden)-1

                logits = self.fc(h_, weights["fc_W"+str(i)], weights["fc_b"+str(i)], relu=False)

        return logits

class LSTMCell(RNNCell):

    '''Vanilla LSTM implemented with same initializations as BN-LSTM'''

    def __init__(self, num_units, W_xh, W_hh, bias):

        self.num_units = num_units

        self.W_xh = W_xh

        self.W_hh = W_hh

        self.bias = bias

    @property

    def state_size(self):

        return (self.num_units, self.num_units)

    @property

    def output_size(self):

        return self.num_units

    def __call__(self, x, state, scope=None):

        with tf.variable_scope(scope or type(self).__name__, reuse=tf.AUTO_REUSE):

            c, h = state

            # hidden = tf.matmul(x, W_xh) + tf.matmul(h, W_hh) + bias

            # improve speed by concat.

            concat = tf.concat([x, h], 1)

            W_both = tf.concat([self.W_xh, self.W_hh], 0)

            hidden = tf.matmul(concat, W_both) + self.bias

            i, j, f, o = tf.split(hidden, 4, axis=1)

            new_c = c * tf.sigmoid(f) + tf.sigmoid(i) * tf.tanh(j)

            new_h = tf.tanh(new_c) * tf.sigmoid(o)

            return new_h, (new_c, new_h)

class CNN(BaseModel):

    """

    Build a convolutional neural network.

    """

    def __init__(self, data_loader, weights_for_finetune, init_std=0.05, freeze_opt=None, is_finetune=False):

        super().__init__()

        self.is_finetune = is_finetune

        self.freeze_opt = freeze_opt

        print ("freeze_opt: ", self.freeze_opt)

        if self.is_finetune:

            self.finetune_weights = weights_for_finetune

            self.learning_rate = 0.00001

            self.batch_size = 64

            self.num_epochs = 30

        else:

            self.learning_rate = 0.1

            self.batch_size = 128

            self.num_epochs = 200

        # training parameters

        self.dir_name =  "cnn"

        self.dropout = 0.6

        self.decay_rate = 0.9

        self.decay_steps = 10000 / self.batch_size

        self.momentum = 0.95

        self.patience = 10

        self.eval_frequency = self.num_epochs

        # Network Parameters

        self.init_std = init_std

        self.n_hidden = 256 # hidden dimensions of embedding

        self.n_hidden_1 = 128

        self.n_hidden_2 = 128

        self.n_words = data_loader.n_words

        self.n_classes = FLAGS.n_classes

        self.n_filters = 128

        self.num_input = data_loader.dim_input

        self.timesteps = data_loader.timesteps

        self.code_size = data_loader.code_size

        self.dim_hidden = [self.n_hidden_1, self.n_hidden_2, FLAGS.n_classes]

        self.filter_sizes = [3, 4, 5]

        self.weights_for_init = dict() # to store the value of learned params

        self.weights_for_finetune = dict()

        print('method', self.dir_name, 'data shape:', self.num_input, 'batch size:', self.batch_size, 'learning rate:', self.learning_rate, \

              'momentum:', self.momentum, 'patience:', self.patience)

        self.build_model()

    # Methods to construct the computational graph

    def build_model(self):

        """Build the computational graph with memory network of the model."""

        self.graph = tf.Graph()

        with self.graph.as_default():

            # Inputs.

            with tf.name_scope('inputs'):

                # tf Graph input

                self.ph_data = tf.placeholder(tf.int32, (self.batch_size, self.timesteps, self.code_size), 'data')

                self.ph_labels = tf.placeholder(tf.int32, (self.batch_size), 'labels')

                self.ph_dropout = tf.placeholder(tf.float32, (), 'dropout')

                self.ph_training = tf.placeholder(tf.bool, name='trainingFlag')

            # Construct model

            op_logits = self._inference(self.ph_data, self.ph_dropout, self.ph_training)

            self.op_loss, self.op_loss_average = self.loss(op_logits)

            self.op_train = self.training(self.op_loss, self.learning_rate,

                    self.decay_steps, self.decay_rate, self.momentum)

            self.op_prediction = self._get_prediction(op_logits)

            # Initialize variables, i.e. weights and biases.

            self.op_init = tf.global_variables_initializer()

            if self.is_finetune is not True:

                self.op_weights = self.get_op_variables()

            else:

                print (tf.trainable_variables())

            # Summaries for TensorBoard and Save for model parameters.

            self.op_summary = tf.summary.merge_all()

            self.op_saver = tf.train.Saver(max_to_keep=5)

        self.graph.finalize()

    def get_op_variables(self):

        op_weights = dict()

        op_var = tf.trainable_variables()

        # embedding

        op_weights["emb_W"] = [v for v in op_var if "emb_W" in v.name][0]

        # cnn

        for i, filter_size in enumerate(self.filter_sizes):

            op_weights["conv_W"+str(filter_size)] = [v for v in op_var if "conv_W"+str(filter_size) in v.name][0]

            op_weights["conv_b"+str(filter_size)] = [v for v in op_var if "conv_b"+str(filter_size) in v.name][0]

        # fully connected

        for i, dim in enumerate(self.dim_hidden):

            op_weights["fc_W"+str(i)] = [v for v in op_var if "fc_W"+str(i) in v.name][0]

            op_weights["fc_b"+str(i)] = [v for v in op_var if "fc_b"+str(i) in v.name][0]

        return op_weights

    def build_emb_weights(self, weights):

        weights["emb_W"] = tf.Variable(tf.random_normal([self.n_words, self.n_hidden], stddev=self.init_std), name="emb_W")

        weights["emb_mask_W"] = tf.get_variable("mask_padding", initializer=MASK_ARRAY, dtype="float32", trainable=False)

        return weights

    def embedding(self, x, Wemb, Wemb_mask):

        _x = tf.nn.embedding_lookup(Wemb, x) # recs size is (batch_size, timesteps, n_words)

        _x_mask = tf.nn.embedding_lookup(Wemb_mask, x)

        emb_vecs = tf.multiply(_x, _x_mask) # broadcast

        emb_vecs = tf.reduce_sum(emb_vecs, 2)

        self.emb_expanded = tf.expand_dims(emb_vecs, -1)

        return emb_vecs

    def build_conv_weights(self, weights):

        for i, filter_size in enumerate(self.filter_sizes):

            filter_shape = [filter_size, self.n_hidden, 1, self.n_filters]

            weights["conv_W"+str(filter_size)] = tf.Variable(tf.truncated_normal(filter_shape, stddev=0.1), name="conv_W"+str(filter_size))

            weights["conv_b"+str(filter_size)] = tf.Variable(tf.constant(0.1, shape=[self.n_filters]), name="conv_b"+str(filter_size))

        return weights

    def conv(self, weights, is_training):

        '''Create a convolution + maxpool layer for each filter size'''

        pooled_outputs = []

        for i, filter_size in enumerate(self.filter_sizes):

            W = weights["conv_W"+str(filter_size)]

            b = weights["conv_b"+str(filter_size)]

            with tf.name_scope("conv-maxpool-%s" % filter_size):

                # Convolution Layer

                conv_ = tf.nn.conv2d(

                    self.emb_expanded,

                    W,

                    strides=[1, 1, 1, 1],

                    padding="VALID",

                    name="conv")

                # Apply nonlinearity

                h = tf.nn.leaky_relu(tf.nn.bias_add(conv_, b), name="relu")

                # h = layers.batch_norm(h, updates_collections=None,

                #                          decay=0.99,

                #                          scale=True, center=True,

                #                          is_training=is_training)

                # Maxpooling over the outputs

                pooled = tf.nn.max_pool(

                h,

                ksize=[1, self.timesteps - filter_size + 1, 1, 1],

                strides=[1, 1, 1, 1],

                padding='VALID',

                name="pool")

                pooled_outputs.append(pooled)

        # Combine all the pooled features

        num_filters_total = self.n_filters * len(self.filter_sizes)

        h_pool = tf.concat(pooled_outputs, 3)

        h_pool_flat = tf.reshape(h_pool, [-1, num_filters_total])

        return h_pool_flat

    # Create model

    def _inference(self, x, dropout, is_training=True):

        with tf.variable_scope('pretrain_model', reuse=None) as training_scope:

            weights = {}

            if self.freeze_opt == None:

                weights = self.build_emb_weights(weights)

                weights = self.build_conv_weights(weights)

                weights = self.build_fc_weights(self.n_filters * len(self.filter_sizes), weights)

                with tf.variable_scope("embedding"):

                    self.embedding(x, weights["emb_W"], weights["emb_mask_W"])

                # convolutional network

                with tf.variable_scope("conv"):

                    hout = self.conv(weights, is_training)

                with tf.variable_scope("dropout"):

                    h_ = layers.dropout(hout, keep_prob=dropout)

                for i, dim in enumerate(self.dim_hidden[:-1]):

                    h_ = self.fc(h_, weights["fc_W"+str(i)], weights["fc_b"+str(i)])

                    h_ = tf.nn.dropout(h_, dropout)

                # Logits linear layer, i.e. softmax without normalization.

                N, Min = h_.get_shape()

                i = len(self.dim_hidden)-1

                logits = self.fc(h_, weights["fc_W"+str(i)], weights["fc_b"+str(i)], relu=False)

            else:

                with tf.variable_scope("embedding"):

                    Wemb = self.finetune_weights["emb_W"]

                    Wemb_mask = tf.get_variable("mask_padding", initializer=MASK_ARRAY, dtype="float32", trainable=False)

                    self.embedding(x, Wemb, Wemb_mask)

                # convolutional network

                with tf.variable_scope("conv"):

                    # w = {}

                    # for i, filter_size in enumerate(self.filter_sizes):

                    #     w["conv_W"+str(filter_size)] = self.finetune_weights["conv_W"+str(filter_size)]

                    #     w["conv_b"+str(filter_size)] = self.finetune_weights["conv_b"+str(filter_size)]

                    hout = self.conv(self.finetune_weights, is_training)

                with tf.variable_scope("dropout"):

                    h_ = layers.dropout(hout, keep_prob=dropout)

                for i, dim in enumerate(self.dim_hidden[:-1]):

                    Wfc = self.finetune_weights["fc_W"+str(i)]

                    bfc = self.finetune_weights["fc_b"+str(i)]

                    h_ = self.fc(h_, Wfc, bfc)

                    h_ = tf.nn.dropout(h_, dropout)

                # finetune the last layer

                i = len(self.dim_hidden)-1

                weights = {}

                dim_in = self.n_hidden_2

                weights["fc_W"+str(i)] = self.weight_variable([int(dim_in), FLAGS.n_classes], name="fc_W"+str(i))

                weights["fc_b"+str(i)] = self.bias_variable([FLAGS.n_classes], name="fc_b"+str(i))

                # Logits linear layer, i.e. softmax without normalization.

                N, Min = h_.get_shape()

                i = len(self.dim_hidden)-1

                logits = self.fc(h_, weights["fc_W"+str(i)], weights["fc_b"+str(i)], relu=False)

        return logits