import os

import numpy as np
import tensorflow as tf
from scipy.stats import ks_2samp
from tensorflow.contrib.rnn import LSTMCell
from tensorflow.contrib.rnn.python.ops import core_rnn_cell
from tensorflow.distributions import Gamma

from rnpp.npp import logger
from .rpp_models import RnnHawkes, Model

_Linear = core_rnn_cell._Linear


class MempoolArrivals(RnnHawkes):

    def __init__(self, args):
        super(MempoolArrivals, self).__init__(args)
        self.sample_dist = self._sample_mark(self.output)

    def log_likelihood(self):
        with tf.name_scope('log_likelihood'):
            with tf.variable_scope("Marks/log_likelihood", reuse=tf.AUTO_REUSE):
                dist_parameters = self._service_distribution_parameters(self.output)
                if self.args.distribution == "Gamma":
                    alpha, beta = tf.split(
                        dist_parameters, 2, self._one, name='distribution_split')
                else:
                    beta = self.dist_parameters
                    alpha = tf.ones_like(beta)
            return self._marks_log_likelihood(self.arrivals, alpha, beta)

    def _build_service_distribution(self):
        with tf.name_scope("build_service_distribution"):
            self.dist_parameters = self._service_distribution_parameters(
                self.output)

    def _marks_log_likelihood(self, x, alpha, beta):
        """
        :param x: [B, T, M]
        :param logits: [B, T, M, K]
        :return: [B, T]
        """
        block_volume_dist = Gamma(alpha, beta)
        shape = tf.shape(x)
        marks = tf.reshape(x[:, :, -1],
                           (shape[0] * shape[1], 1))
        bl_log_like = block_volume_dist.log_prob(tf.abs(marks) + 1e-12)
        bl_log_like = tf.reshape(bl_log_like, (shape[0], shape[1]))

        return bl_log_like

    def _service_distribution_parameters(self, h_a):
        """
        Args:
            output: [B, T, H]
            covariates: [B, T, D-3]
        """
        with tf.variable_scope(
                "Marks/distribution_parameters", reuse=tf.AUTO_REUSE):
            output = tf.reshape(
                h_a, [-1, self.args.hidden_units])

            cell_initializer = None

            if self.args.initializer == 'constant':
                cell_initializer = tf.initializers.constant(
                    self.args.init_value)

            marks_layer = _Linear(
                output, self.args.marks_layer_size, True, kernel_initializer=cell_initializer)
            z = tf.nn.tanh(marks_layer(output))

            with tf.variable_scope("dist_parameters", reuse=tf.AUTO_REUSE):
                if self.args.distribution == "Gamma":
                    num_dist_params = 2
                if self.args.distribution == "Exp":
                    num_dist_params = 1
                distribution_layer = _Linear(
                    z, num_dist_params, True, kernel_initializer=cell_initializer)
                p = distribution_layer(z)

                return tf.exp(p)

    def _sample_mark(self, output):
        shape = tf.shape(output)
        dist_parameters = self._service_distribution_parameters(
            output)
        if self.args.distribution == "Gamma":
            alpha, beta = tf.split(
                dist_parameters, 2, self._one, name='distribution_split')
        else:
            beta = tf.cast(dist_parameters, dtype=tf.float64)
            alpha = tf.ones_like(beta, dtype=tf.float64)

        dist = Gamma(alpha, beta)
        samples = tf.transpose(dist.sample(self.args.n_mc_samples)[
                               :, :, 0])  # [B*T, n_mc_samples]

        # [B, T, n_mc_samples]
        samples = tf.reshape(
            samples, (shape[0], shape[1], self.args.n_mc_samples))
        return samples  # [B, T, n_mc_samples]

    def _predict_mark(self, output):
        """
        Args:
            output:
            covariates:
        """
        shape = tf.shape(output)
        dist_parameters = self._service_distribution_parameters(
            output)
        if self.args.distribution == "Gamma":
            alpha, beta = tf.split(
                dist_parameters, 2, self._one, name='distribution_split')
        else:
            alpha = 1.
            beta = dist_parameters
        rop = tf.divide(alpha, beta)
        rop = tf.reshape(rop, (shape[0], shape[1]))
        return rop

    def _build_one_step_rmse(self):
        marks_real = self.arrivals[:, :, -1]

        marks_predicted = self._predict_mark(self.output)
        marks_predicted = tf.multiply(self.mask, marks_predicted)
        se = tf.square(tf.subtract(marks_predicted, marks_real))

        self.one_step_rmse = se

    def train(self, train_set, test_set, sess):
        """

        Args:
            train_data:
            test_data:
            seq_len:
            mask:
            sess:

        Returns:

        """
        train_data, train_seq_len, train_mask = train_set
        test_data, test_seq_len, test_mask = test_set
        sess.run(self.initialize)
        summary_writer = self.summary_writer(self.tensorboard_dir, sess)
        NUM_BATCHES, _, NUM_SEQ, _, _ = train_data.shape
        TOTAL_ITERATIONS = self.args.num_epochs * NUM_BATCHES * (NUM_SEQ - 1.)
        num_epochs = self.args.num_epochs
        try:
            for epoch in range(num_epochs):
                state = sess.run(self.reset_state, {
                    self.arrivals: train_data[0, :, 0]})
                costs = []
                for batch_ix in range(NUM_BATCHES):
                    for seq_ix in range(NUM_SEQ):
                        x = train_data[batch_ix, :, seq_ix]
                        s_l = train_seq_len[batch_ix, :, seq_ix]
                        m = train_mask[batch_ix, :, seq_ix]
                        feed_dict = {
                            self.arrivals: x,
                            self.state: state,
                            self.seq_len: s_l,
                            self.mask: m
                        }

                        fetches = [
                            self.cost, self.next_state, self.global_step
                        ]

                        if seq_ix > 0:
                            fetches.append(self.train_step)
                            cost, state, iteration, _ = sess.run(
                                fetches, feed_dict)
                        else:
                            cost, state, iteration = sess.run(
                                fetches, feed_dict)

                        if self.logging_interval(iteration,
                                                 self.args.logging_interval):
                            logger.info(
                                "Epoch {0} ({1:.2%}) Step {2}: {3:4f}".format(
                                    epoch, iteration / TOTAL_ITERATIONS,
                                    iteration, cost))
                        costs.append(cost)
                N, T_NUM_SEQ, _, _ = test_data.shape
                se = []
                predictions = []
                state = state[:N]
                fetches = [self.one_step_rmse, self.next_state, self.sample_dist]
                for t_seq_ix in range(T_NUM_SEQ):
                    se_seq, state, prediction = sess.run(
                        fetches, {
                            self.arrivals: test_data[:, t_seq_ix],
                            self.state: state,
                            self.seq_len: test_seq_len[:, t_seq_ix],
                            self.mask: test_mask[:, t_seq_ix]
                        })

                    prediction = np.concatenate(
                        (prediction, test_data[:, t_seq_ix]), axis=2)  # [B, T, MC + 1]
                    predictions.append(prediction)
                    se.append(se_seq)
                predictions = np.concatenate(predictions, axis=1)
                dist_mean = np.mean(
                    predictions[:, :, :self.args.n_mc_samples])
                rmse = np.mean(se[1:])
                service_times = np.asarray(predictions)
                ks_d, p = ks_2samp(
                    service_times[:, :, :-1].flatten(), service_times[:, :, -1].flatten())
                logger.info("RMSE on test set epoch {1}: {0:4f}, dist_mean: {2:4f}, d: {3:4f}, p: {4:4f}".format(
                    rmse, epoch, dist_mean, ks_d, p))
                if epoch % 10 == 0:
                    np.save(str(os.path.join(self.checkpoint_dir, "sample-iter-{}".format(epoch))), service_times)
                if ks_d < self.best_model_loss:
                    self.best_model_loss = ks_d
                    self._save_best_model(sess)

                self.store_training_epoch(sess, summary_writer, epoch,
                                          iteration, np.mean(costs), (rmse, ks_d))

        except KeyboardInterrupt:
            logger.info("Stop training at epoch {0}, iteration {1}".format(
                epoch, iteration))
        summary_writer.close()


class MempoolService(RnnHawkes):

    def __init__(self, args):
        super(MempoolService, self).__init__(args)
        state = tf.concat((self.output, self.arrivals_hidden), -1)
        self.sample_dist = self._sample_mark(state)

    def _build_placeholders(self):
        super(MempoolService, self)._build_placeholders()
        with tf.name_scope('Inputs'):
            self.arrivals_hidden = tf.placeholder(tf.float64, [None, None, self.args.arrivals_hidden_units])

    def _build_recurrent(self):
        cell_initializer = tf.initializers.random_uniform()
        if self.args.initializer == 'constant':
            cell_initializer = tf.initializers.constant(self.args.init_value)
        with tf.variable_scope(
                'PointProcess/RNN', initializer=cell_initializer):
            if self.args.cell == 'Basic':
                self._cell = tf.nn.rnn_cell.BasicRNNCell(self.args.hidden_units, tf.nn.relu)
            else:
                self._cell = tf.nn.rnn_cell.GRUCell(self.args.hidden_units, tf.nn.relu)
            self.reset_state = self._cell.zero_state(
                tf.shape(self.arrivals)[0], dtype=tf.float64)
            self.reset_state += tf.random_normal(
                tf.shape(self.reset_state), dtype=tf.float64)
            self.state = tf.placeholder(tf.float64,
                                        self.reset_state.get_shape(), "state")
            self.output, self.next_state = self._get_hidden_states(
                self.arrivals[:, :, :3], self.state, self.seq_len)

    def _build_one_step_ahead_prediction(self):
        self.one_step_ahead_prediction_is = self._predict_one_step_ahead_inverse_sampling(
            self.arrivals[:, :, :3], self.state)

    def log_likelihood(self):
        with tf.name_scope('log_likelihood'):
            hidden = tf.concat((self.output, self.arrivals_hidden), -1)
            with tf.variable_scope("PointProcess/Marks/log_likelihood", reuse=tf.AUTO_REUSE):
                dist_parameters = self._service_distribution_parameters(hidden)
                if self.args.distribution == "Gamma":
                    alpha, beta = tf.split(
                        dist_parameters, 2, self._one, name='distribution_split')
                else:
                    beta = dist_parameters
                    alpha = tf.ones_like(beta)
                marks_log_likelihood = self._marks_log_likelihood(self.arrivals, alpha, beta)
            return marks_log_likelihood + self._arrivals_log_likelihood(self.arrivals, hidden)

    def _build_service_distribution(self):
        with tf.name_scope("build_service_distribution"):
            states = tf.concat((self.output, self.arrivals_hidden), axis=-1)
            self.dist_parameters = self._service_distribution_parameters(states)

    def _marks_log_likelihood(self, x, alpha, beta):
        """
        :param x: [B, T, M]
        :param logits: [B, T, M, K]
        :return: [B, T]
        """
        block_volume_dist = Gamma(alpha, beta)
        shape = tf.shape(x)
        marks = tf.reshape(x[:, :, -1],
                           (shape[0] * shape[1], 1))
        bl_log_like = block_volume_dist.log_prob(tf.abs(marks) + 1e-12)
        bl_log_like = tf.reshape(bl_log_like, (shape[0], shape[1]))

        return bl_log_like

    def _service_distribution_parameters(self, h_a):
        """
        Args:
            output: [B, T, H]
            covariates: [B, T, D-3]
        """
        with tf.variable_scope(
                "Marks/distribution_parameters", reuse=tf.AUTO_REUSE):
            output = tf.reshape(
                h_a, [-1, self.args.hidden_units + self.args.arrivals_hidden_units])

            cell_initializer = None

            if self.args.initializer == 'constant':
                cell_initializer = tf.initializers.constant(
                    self.args.init_value)

            marks_layer = _Linear(
                output, self.args.marks_layer_size, True, kernel_initializer=cell_initializer)
            z = tf.nn.tanh(marks_layer(output))

            with tf.variable_scope("dist_parameters", reuse=tf.AUTO_REUSE):
                if self.args.distribution == "Gamma":
                    num_dist_params = 2
                if self.args.distribution == "Exp":
                    num_dist_params = 1
                distribution_layer = _Linear(
                    z, num_dist_params, True, kernel_initializer=cell_initializer)
                p = distribution_layer(z)

                return tf.exp(p)

    def _sample_mark(self, output):
        shape = tf.shape(output)
        dist_parameters = self._service_distribution_parameters(
            output)
        if self.args.distribution == "Gamma":
            alpha, beta = tf.split(
                dist_parameters, 2, self._one, name='distribution_split')
        else:
            beta = tf.cast(dist_parameters, dtype=tf.float64)
            alpha = tf.ones_like(beta, dtype=tf.float64)

        dist = Gamma(alpha, beta)
        samples = tf.transpose(dist.sample(self.args.n_mc_samples)[
                               :, :, 0])  # [B*T, n_mc_samples]

        # [B, T, n_mc_samples]
        samples = tf.reshape(
            samples, (shape[0], shape[1], self.args.n_mc_samples))
        return samples  # [B, T, n_mc_samples]

    def _predict_mark(self, output):
        """
        Args:
            output:
            covariates:
        """
        shape = tf.shape(output)
        dist_parameters = self._service_distribution_parameters(
            output)
        if self.args.distribution == "Gamma":
            alpha, beta = tf.split(
                dist_parameters, 2, self._one, name='distribution_split')
        else:
            alpha = 1.
            beta = dist_parameters
        rop = tf.divide(alpha, beta)
        rop = tf.reshape(rop, (shape[0], shape[1]))
        return rop

    def _predict_one_step_ahead_inverse_sampling(self, x, last_state):
        """

        Args:
            x:
            last_state:

        Returns:

        """
        with tf.name_scope('inverse_sampling'):
            states, _ = self._get_hidden_states(x, last_state, None)
            x_shape = tf.shape(x)
            y = tf.random_uniform([
                x_shape[0], x_shape[1],
                tf.cast(self.args.inverse_sampling_size, tf.int32)
            ])
            states = tf.concat((states, self.arrivals_hidden), -1)
            alpha = self._past_influence(states) + self.b_t

            alpha = tf.expand_dims(alpha, 2)
            log_inv_sample = tf.cast(
                tf.log(tf.divide(-1., y - 1.)), tf.float64)
            inv_sample = tf.divide(
                tf.log(
                    tf.abs(
                        tf.multiply(
                            self.w_t, log_inv_sample + tf.divide(
                                tf.exp(alpha), self.w_t)))) - alpha, self.w_t)
            return x[:, :, 0] + tf.reduce_mean(inv_sample, axis=2)

    def _build_one_step_rmse(self):
        y_predicted = self.one_step_ahead_prediction_is
        y_actual = self.arrivals[:, :, -2]

        y_predicted = tf.multiply(y_predicted, self.mask)
        se_arrival = tf.square(tf.subtract(y_predicted, y_actual))

        marks_real = self.arrivals[:, :, -1]
        states = tf.concat((self.output, self.arrivals_hidden), -1)
        marks_predicted = self._predict_mark(states)
        marks_predicted = tf.multiply(self.mask, marks_predicted)
        se_mark = tf.square(tf.subtract(marks_predicted, marks_real))

        self.one_step_rmse = tf.tuple((se_arrival, se_mark))

    def intensity(self):
        with tf.name_scope('intensity'):
            hidden_states = tf.concat((self.output, self.arrivals_hidden), axis=-1)
            return self._intensity(self.arrivals, hidden_states)

    def __log_intensity(self, x, hidden_states):
        """
        Args:
            x: [B, T, D]
            hidden_states: [B, T, H]

        Returns:
            [B, T]
        """
        with tf.name_scope('log_intensity'):
            past_influence = self._past_influence(hidden_states)
            current_influence = tf.multiply(self.w_t, x[:, :, -1] - x[:, :, 0],
                                            'current_influence')
            return past_influence + current_influence + self.b_t

    def train(self, train_set, test_set, hidden_train, hidden_test, sess):
        """

        Args:
            train_data:
            test_data:
            seq_len:
            mask:
            sess:

        Returns:

        """
        train_data, train_seq_len, train_mask = train_set
        test_data, test_seq_len, test_mask = test_set
        sess.run(self.initialize)
        summary_writer = self.summary_writer(self.tensorboard_dir, sess)
        NUM_BATCHES, _, NUM_SEQ, _, _ = train_data.shape
        TOTAL_ITERATIONS = self.args.num_epochs * NUM_BATCHES * (NUM_SEQ - 1.)
        num_epochs = self.args.num_epochs
        try:
            for epoch in range(num_epochs):
                state = sess.run(self.reset_state, {
                    self.arrivals: train_data[0, :, 0]})
                for batch_ix in range(NUM_BATCHES):
                    costs = []
                    for seq_ix in range(NUM_SEQ):
                        x = train_data[batch_ix, :, seq_ix]
                        s_l = train_seq_len[batch_ix, :, seq_ix]
                        m = train_mask[batch_ix, :, seq_ix]
                        h = hidden_train[batch_ix, :, seq_ix]

                        feed_dict = {
                            self.arrivals: x,
                            self.state: state,
                            self.seq_len: s_l,
                            self.mask: m,
                            self.arrivals_hidden: h
                        }

                        fetches = [
                            self.cost, self.next_state, self.global_step
                        ]

                        if seq_ix > 0:
                            fetches.append(self.train_step)
                            cost, state, iteration, _ = sess.run(
                                fetches, feed_dict)
                        else:
                            cost, state, iteration = sess.run(
                                fetches, feed_dict)

                        if self.logging_interval(iteration,
                                                 self.args.logging_interval):
                            logger.info(
                                "Epoch {0} ({1:.2%}) Step {2}: {3:4f}".format(
                                    epoch, iteration / TOTAL_ITERATIONS,
                                    iteration, cost))
                        costs.append(cost)
                N, T_NUM_SEQ, _, _ = test_data.shape
                se_arr = []
                se_marks = []
                predictions_marks = []
                state = sess.run(self.reset_state, {
                    self.arrivals: test_data[:, 0]})
                state = state[:N]

                fetches = [self.one_step_rmse, self.next_state, self.sample_dist]
                for t_seq_ix in range(T_NUM_SEQ):
                    se_seq, state, prediction = sess.run(
                        fetches, {
                            self.arrivals: test_data[:, t_seq_ix],
                            self.state: state,
                            self.seq_len: test_seq_len[:, t_seq_ix],
                            self.mask: test_mask[:, t_seq_ix],
                            self.arrivals_hidden: hidden_test[:, t_seq_ix]
                        })
                    prediction = np.concatenate(
                        (prediction, test_data[:, t_seq_ix]), axis=2)  # [B, T, MC + 1]
                    predictions_marks.append(prediction)
                    se_arr.append(se_seq[0])
                    se_marks.append(se_seq[1])
                predictions_marks = np.concatenate(predictions_marks, axis=1)
                dist_mean = np.mean(
                    predictions_marks[:, :, :self.args.n_mc_samples])
                rmse_arr = np.sqrt(np.mean(se_arr[1:]))
                rmse_mark = np.sqrt(np.mean(se_marks[1:]))
                predictions_marks = np.asarray(predictions_marks)
                ks_d, p = ks_2samp(
                    predictions_marks[:, :, :-1].flatten(), predictions_marks[:, :, -1].flatten())
                logger.info("RMSE on test set epoch {1}: {0:4f}, dist_mean: {2:4f}, d: {3:4f}, p: {4:4f}".format(
                    rmse_arr, epoch, dist_mean, ks_d, p))

                if epoch % 10 == 0:
                    np.save(str(os.path.join(self.checkpoint_dir, "sample-iter-{}".format(epoch))), predictions_marks)
                if ks_d < self.best_model_loss:
                    self.best_model_loss = ks_d
                    self._save_best_model(sess)

                self.store_training_epoch(sess, summary_writer, epoch,
                                          iteration, np.sum(costs), (rmse_arr, ks_d))

        except KeyboardInterrupt:
            logger.info("Stop training at epoch {0}, iteration {1}".format(
                epoch, iteration))
        summary_writer.close()


class RecurrentAdversarialTransactions(Model):

    def __init__(self, args):
        self.model = "RecurrentAdversarialTransactions"
        super(RecurrentAdversarialTransactions, self).__init__(args)
        self._build_placeholders()
        self._build_service_distribution()
        self.__build_training()
        self.__build_train_step()
        self._build_one_step_rmse()
        self.initialize = tf.global_variables_initializer()
        self.summary = tf.summary.merge_all()
        self._configure_logger(self.args.log.upper(),
                               "%(asctime)s - %(levelname)s - %(message)s")
        self.saver = tf.train.Saver(max_to_keep=100)

        self.sample_dist = self._get_predicted_marks()

    def _build_recurrent(self):
        pass

    def _build_one_step_ahead_prediction(self):
        pass

    def _predict_one_step_ahead(self, x, last_state):
        pass

    def _arrivals_log_likelihood(self, x, hidden_state):
        pass

    def inference(self, x, sess):
        pass

    def __build_training(self):
        with tf.name_scope("Train"):
            self.rmse_arrivals = tf.Variable(
                float("inf"), False, dtype=tf.float64, name="rmse_arrivals")
            self.rmse_marks = tf.Variable(
                float("inf"), False, dtype=tf.float64, name="rmse_marks")
            self.loss = tf.Variable(
                float("inf"), False, dtype=tf.float64, name="loss")
            self.loss_discriminator = tf.Variable(
                float("inf"), False, dtype=tf.float64, name="loss_dis")

            tf.summary.scalar(self.rmse_arrivals.op.name, self.rmse_arrivals)
            tf.summary.scalar(self.rmse_marks.op.name, self.rmse_marks)
            tf.summary.scalar(self.loss.op.name, self.loss)
            tf.summary.scalar(self.loss_discriminator.op.name,
                              self.loss_discriminator)

            with tf.name_scope("Losses"):
                self.mark_loss = self._get_loss()
                self.dis_loss = self._get_discriminator_loss()

            self.__build_train_step()

    def __build_train_step(self):

        optimizer = tf.train.AdamOptimizer(
            learning_rate=self.args.learning_rate_gan,
            beta1=0.5,
            beta2=0.9)

        # Trainable variables within the GENERATOR

        train_var = tf.get_collection(
            tf.GraphKeys.GLOBAL_VARIABLES, scope="GAN/generator")
        grad = tf.gradients(self.mark_loss, train_var)
        self.train_step = optimizer.apply_gradients(
            zip(grad, train_var),
            global_step=self.global_step,
            name="train_step_gen")

        # Trainable variables within the DISCRIMINATOR

        train_dis_var = tf.get_collection(
            tf.GraphKeys.GLOBAL_VARIABLES, scope="GAN/discriminator")
        grad_dis = tf.gradients(self.dis_loss, train_dis_var)
        self.train_step_dis = optimizer.apply_gradients(
            zip(grad_dis, train_dis_var), name="train_step_dis")

    def _build_placeholders(self):
        self.arrivals = tf.placeholder(
            tf.float64, [None, self.args.num_steps, self.args.input_dim], 'input_data')

        self.seq_len = tf.placeholder(tf.int32, [None], 'seq_len')

        self.mask = tf.placeholder(tf.float64, [None, self.args.num_steps], 'mask')

        self.dropout_prob = tf.placeholder(
            tf.float64, shape=[], name="keep_prob_dropout")

        self.state = self.__get_rnn_state_placeholder(
            self.args.gen_dim, name="state_generator")

        self.state_dis_real = self.__get_rnn_state_placeholder(
            self.args.dis_dim, name="state_discriminator_real")

        self.state_dis_fake = self.__get_rnn_state_placeholder(
            self.args.dis_dim, name="state_discriminator_fake")

        self.state_dis_interpol = self.__get_rnn_state_placeholder(
            self.args.dis_dim, name="state_discriminator_interpol")

    @staticmethod
    def __get_rnn_state_placeholder(h_dim, name=None):

        c = tf.placeholder(tf.float64, shape=[None, h_dim], name=name)
        h = tf.placeholder(tf.float64, shape=[None, h_dim], name=name)

        return tf.nn.rnn_cell.LSTMStateTuple(c, h)

    def _build_service_distribution(self):
        self._build_recurrent_adversarial()

    def _build_recurrent_adversarial(self):
        covariates = self.arrivals[:, :, :-1]
        u_real = tf.reshape(self.arrivals[:, :, -1],
                            [-1, 1])  # target number of unconfirmed transactions [B*T, 1]

        with tf.name_scope("build_generator-discriminator"):
            # Generator output
            self._u_gen, self.next_state = self._generator(
                covariates, self.state)  # [B*T, 1], [B, T, H]

            self._logit_real, self.next_state_dis_real = self._discriminator(
                covariates, u_real, self.state_dis_real)  # [B*T, 1]

            self._logit_fake, self.next_state_dis_fake = self._discriminator(
                covariates, self._u_gen, self.state_dis_fake)  # [B*T, 1]

            self._grad_penalty, self.next_state_dis_interpol = self._gradient_penalty_term(
                u_real, self._u_gen, covariates)

    def _generator(self, covariates, state):
        """
        GAN generator
        covariates [B, T, D-1]
        """

        with tf.variable_scope("GAN/generator", reuse=tf.AUTO_REUSE):
            with tf.variable_scope("hidden_layers", reuse=tf.AUTO_REUSE):
                tau = tf.reshape(covariates[:, :, 0], [-1, 1])
                u = tf.reshape(covariates[:, :, 1], [-1, 1])

                # embedding
                cell_initializer = tf.contrib.layers.xavier_initializer()
                tau = tf.layers.dense(
                    tau,
                    self.args.gen_dim,
                    kernel_initializer=cell_initializer,
                    name="layer_hidden_tau")

                u = tf.layers.dense(
                    u,
                    self.args.gen_dim,
                    kernel_initializer=cell_initializer,
                    name="layer_hidden_u")

                # sampling prior normal distr.
                z = tf.random_normal(tf.shape(tau), 0.0, 1.0, dtype=tf.float64)
                z = tf.layers.dense(
                    z,
                    self.args.gen_dim,
                    use_bias=False,
                    kernel_initializer=cell_initializer,
                    name="layer_noise_0")

                # hidden state
                h = tf.nn.leaky_relu(tau + u + z)  # [B*T, D_GAN]
                h = tf.reshape(h, [-1, self.args.num_steps, self.args.gen_dim])
                h = tf.transpose(h, [1, 0, 2])
                h = tf.unstack(h)

                # recurrent calculation
                cell = LSTMCell(
                    self.args.gen_dim,
                    state_is_tuple=True,
                    reuse=tf.AUTO_REUSE)
                h, next_state = tf.nn.static_rnn(
                    cell=cell,
                    inputs=h,
                    sequence_length=self.seq_len,
                    initial_state=state)

                h = tf.stack(h)
                h = tf.transpose(h, [1, 0, 2])

                self.hidden_states_gen = h

                h = tf.reshape(h, [-1, self.args.gen_dim])

                if self.args.n_layers_gen > 1:
                    for i in range(self.args.n_layers_gen - 1):
                        h = tf.layers.dense(
                            h,
                            self.args.gen_dim,
                            kernel_initializer=cell_initializer,
                            name="layer_hidden_{}".format(i + 1))
                        if self.args.noisy_layers:
                            z = tf.random_normal(
                                tf.shape(h), 0.0, 1.0, dtype=tf.float64)
                            z = tf.layers.dense(
                                z,
                                self.args.gen_dim,
                                use_bias=False,
                                kernel_initializer=cell_initializer,
                                name="layer_noise_{}".format(i + 1))
                            h = h + z
                        if self.args.dropout:
                            h = tf.nn.dropout(h, self.dropout_prob)
                        h = tf.nn.leaky_relu(h)  # [B*T, D_GAN]

            with tf.variable_scope("generator_output", reuse=tf.AUTO_REUSE):
                h = tf.layers.dense(h, 1, kernel_initializer=cell_initializer)

        return tf.exp(h), next_state  # [B*T, 1], [B, T, h_dim]

    def _discriminator(self, covariates, u_new, state):
        """
        GAN discriminator
        h (arrival hidden state) [B, T, H]
        covariates [B, T, D-1]
        s (unconfirmed transactions) [B*T, 1]
        """
        with tf.variable_scope("GAN/discriminator", reuse=tf.AUTO_REUSE):
            with tf.variable_scope("hidden_layers", reuse=tf.AUTO_REUSE):
                tau = tf.reshape(covariates[:, :, 0], [-1, 1])
                u_old = tf.reshape(covariates[:, :, 1], [-1, 1])
                u_new = tf.reshape(u_new, [-1, 1])

                # embedding
                cell_initializer = tf.contrib.layers.xavier_initializer()
                tau = tf.layers.dense(
                    tau,
                    self.args.dis_dim,
                    kernel_initializer=cell_initializer,
                    name="layer_hidden_tau")

                u_old = tf.layers.dense(
                    u_old,
                    self.args.dis_dim,
                    kernel_initializer=cell_initializer,
                    name="layer_hidden_u_old")

                u_new = tf.layers.dense(
                    u_new,
                    self.args.dis_dim,
                    use_bias=False,
                    kernel_initializer=cell_initializer,
                    name="layer_hidden_u_new")

                # hidden state
                h = tf.nn.leaky_relu(tau + u_old + u_new)  # [B*T, D_GAN]
                h = tf.reshape(h, [-1, self.args.num_steps, self.args.gen_dim])
                h = tf.transpose(h, [1, 0, 2])
                h = tf.unstack(h)

                # recurrent calculation
                cell = LSTMCell(
                    self.args.dis_dim,
                    state_is_tuple=True,
                    reuse=tf.AUTO_REUSE)
                h, next_state = tf.nn.static_rnn(
                    cell=cell,
                    inputs=h,
                    sequence_length=self.seq_len,
                    initial_state=state)

                h = tf.stack(h)
                h = tf.transpose(h, [1, 0, 2])
                h = tf.reshape(h, [-1, self.args.dis_dim])

                if self.args.n_layers_dis > 1:
                    for i in range(self.args.n_layers_dis - 1):
                        h = tf.layers.dense(
                            h,
                            self.args.dis_dim,
                            kernel_initializer=cell_initializer,
                            name="layer_hidden_{}".format(i + 1))
                        if self.args.dropout:
                            h = tf.nn.dropout(h, self.dropout_prob)
                        h = tf.nn.leaky_relu(h)  # [B*T, D_GAN]

            with tf.variable_scope(
                    "discriminator_output", reuse=tf.AUTO_REUSE):
                h = tf.layers.dense(h, 1, kernel_initializer=cell_initializer)

        return h, next_state  # [B*T, 1], [B, T, h_dim]

    def _gradient_penalty_term(self, u_real, u_gen, covariates):

        with tf.name_scope("gradient_penalty_term"):
            eta = tf.random_uniform(
                tf.shape(u_real), minval=0.0, maxval=0.1, dtype=tf.float64)

            u_interpol = eta * u_real + (1 - eta) * u_gen
            logit_interpol, next_state_interpol = self._discriminator(covariates,
                                                                      u_interpol, self.state_dis_interpol)

            grad_penalty = tf.gradients(logit_interpol, u_interpol)[0]  # tf.gradients returns a list

            return tf.abs(grad_penalty), next_state_interpol

    def _get_discriminator_loss(self):
        with tf.name_scope("adversarial_loss"):
            with tf.name_scope("discriminator_component"):
                shape = tf.shape(self.arrivals)
                batch_size = shape[0]
                seq_length = shape[1]

                dis_loss = -(self._logit_real - self._logit_fake)

                with tf.name_scope("gradient_penalty"):
                    dis_loss = dis_loss + self.args.lambda_grad_penalty * \
                               tf.square(tf.maximum(tf.cast(0.0, tf.float64),
                                                    self._grad_penalty - 1.0))

                dis_loss = tf.reshape(dis_loss, [batch_size, seq_length])
                dis_loss = tf.multiply(dis_loss, self.mask)
                dis_loss = tf.reduce_sum(dis_loss)

            return dis_loss

    def _get_loss(self):
        with tf.name_scope("adversarial_loss"):
            with tf.name_scope("generator_component"):
                shape = tf.shape(self.arrivals)
                batch_size = shape[0]
                seq_length = shape[1]
                # standard generator loss component of GAN

                gen_loss = -self._logit_fake

                gen_loss = tf.reshape(gen_loss, [batch_size, seq_length])
                gen_loss = tf.multiply(gen_loss, self.mask)
                gen_loss = tf.reduce_sum(gen_loss)

            return gen_loss

    def _predict_mark(self, covariates):
        """
        Args:
            output:
            covariates:
        """

        shape = tf.shape(self.arrivals)
        next_mark, next_state = self._generator(covariates, self.state)
        next_mark = tf.reshape(next_mark, (shape[0], shape[1]))

        return next_mark, next_state

    def _get_predicted_marks(self):

        covariates = self.arrivals[:, :, :-1]

        shape = tf.shape(self.arrivals)

        next_mark, _ = self._generator(covariates,
                                       self.state)

        next_mark = tf.reshape(next_mark, (shape[0], shape[1]))

        next_mark = tf.multiply(self.mask, next_mark)

        next_mark = tf.split(next_mark, self.args.batch_size_test, axis=0)
        next_mark = tf.stack(next_mark)  # [B, MC, T]
        next_mark = tf.transpose(next_mark, [0, 2, 1])

        return next_mark

    def _build_one_step_rmse(self):

        marks_real = self.arrivals[:, :, -1]
        covariates = self.arrivals[:, :, :-1]

        marks_predicted, _ = self._predict_mark(covariates)

        marks_predicted = tf.multiply(self.mask, marks_predicted)

        # Averaged values:
        # 1.- predictions:
        marks_predicted = tf.split(marks_predicted,
                                   self.args.batch_size_test,
                                   axis=0)  # [B, MC, T]
        marks_predicted = tf.stack(marks_predicted)
        marks_predicted = tf.transpose(marks_predicted, [0, 2, 1])  # [B, T, MC]
        marks_predicted = tf.reduce_mean(marks_predicted, axis=2)
        # 2.- real:
        marks_real = tf.split(marks_real,
                              self.args.batch_size_test,
                              axis=0)  # [B, MC, T]
        marks_real = tf.stack(marks_real)
        marks_real = tf.transpose(marks_real, [0, 2, 1])  # [B, T, MC]
        marks_real = tf.reduce_mean(marks_real, axis=2)

        error = tf.abs(marks_predicted - marks_real)
        self.one_step_rmse = error

    def extract_hidden_states(self, train_set, test_set, sess):
        train_data, train_seq_len, _ = train_set
        test_data, test_seq_len, test_mask = test_set
        NUM_BATCHES, B, NUM_SEQ, _, _ = train_data.shape

        test_states = []
        train_states = []

        # Initial state:

        for batch_ix in range(NUM_BATCHES):
            states = []
            state = self._get_zero_state(B, self.args.gen_dim)
            for seq_ix in range(NUM_SEQ):
                x = train_data[batch_ix, :, seq_ix]
                s_l = train_seq_len[batch_ix, :, seq_ix]

                h, state = sess.run([self.hidden_states_gen, self.next_state], {
                    self.arrivals: x,
                    self.state: state,
                    self.seq_len: s_l
                })
                states.append(np.expand_dims(h, 1))
            train_states.append(np.concatenate(states, 1))

        B_T, T_NUM_SEQ, _, _ = test_data.shape
        state = self._get_zero_state(B_T, self.args.gen_dim)
        for seq_ix in range(T_NUM_SEQ):
            h, state = sess.run([self.hidden_states_gen, self.next_state], {
                    self.arrivals: test_data[:, seq_ix],
                    self.state: state,
                    self.seq_len: test_seq_len[:, seq_ix]
                })
            test_states.append(np.expand_dims(h, 1))

        train_states = np.stack(train_states, 0)
        test_states = np.concatenate(test_states, 1)

        return train_states, test_states

    def train(self, train_set, test_set, sess):
        train_data, train_seq_len, train_mask = train_set
        test_data, test_seq_len, test_mask = test_set
        sess.run(self.initialize)
        summary_writer = self.summary_writer(self.tensorboard_dir, sess)
        NUM_BATCHES, B, NUM_SEQ, _, _ = train_data.shape
        TOTAL_ITERATIONS = self.args.num_epochs * NUM_BATCHES * (NUM_SEQ - 1.)
        num_epochs = self.args.num_epochs
        print "train data shape: ", train_data.shape, "; test data shape: ", test_data.shape
        print "train mask shape: ", train_mask.shape, "; test mask shape: ", test_mask.shape
        try:
            for epoch in range(num_epochs):
                mark_loss_set = []
                dis_loss_set = []
                for batch_ix in range(NUM_BATCHES):
                    # Initial state:
                    state = self._get_zero_state(B, self.args.gen_dim)
                    state_dis_real = self._get_zero_state(B, self.args.dis_dim)
                    state_dis_fake = self._get_zero_state(B, self.args.dis_dim)
                    state_dis_interpol = self._get_zero_state(B, self.args.dis_dim)
                    for seq_ix in range(NUM_SEQ):
                        x = train_data[batch_ix, :, seq_ix]
                        s_l = train_seq_len[batch_ix, :, seq_ix]
                        m = train_mask[batch_ix, :, seq_ix]
                        feed_dict = {
                            self.arrivals: x,
                            self.state: state,
                            self.state_dis_real: state_dis_real,
                            self.state_dis_fake: state_dis_fake,
                            self.state_dis_interpol: state_dis_interpol,
                            self.seq_len: s_l,
                            self.mask: m,
                            self.dropout_prob: self.args.dropout_prob
                        }

                        dis_loss, _, state_dis_fake, state_dis_real, state_dis_interpol = sess.run(
                            [self.dis_loss,
                             self.train_step_dis,
                             self.next_state_dis_fake,
                             self.next_state_dis_real,
                             self.next_state_dis_interpol], feed_dict)

                        mark_loss, _, iteration, state = sess.run([self.mark_loss, self.train_step,
                                                                   self.global_step, self.next_state],
                                                                  feed_dict)

                        if self.logging_interval(iteration, self.args.logging_interval):
                            logger.info(
                                "Epoch {0} ({1:.2%}) Step {2}: {3:4f}, {4:4f}"
                                    .format(epoch,
                                            iteration / TOTAL_ITERATIONS,
                                            iteration, mark_loss, dis_loss))

                        mark_loss_set.append(mark_loss)
                        dis_loss_set.append(dis_loss)

                    if self.args.n_discriminator_updates > 1:
                        for i in range(self.args.n_discriminator_updates - 1):
                            state = self._get_zero_state(B, self.args.gen_dim)
                            state_dis_real = self._get_zero_state(
                                B, self.args.dis_dim)
                            state_dis_fake = self._get_zero_state(
                                B, self.args.dis_dim)
                            state_dis_interpol = self._get_zero_state(
                                B, self.args.dis_dim)
                            for seq_ix in range(NUM_SEQ):
                                x = train_data[batch_ix, :, seq_ix]
                                s_l = train_seq_len[batch_ix, :, seq_ix]
                                m = train_mask[batch_ix, :, seq_ix]
                                feed_dict.update({self.arrivals: x,
                                                  self.mask: m,
                                                  self.seq_len: s_l,
                                                  self.state: state,
                                                  self.state_dis_fake: state_dis_fake,
                                                  self.state_dis_real: state_dis_real,
                                                  self.state_dis_interpol: state_dis_interpol,
                                                  self.dropout_prob: self.args.dropout_prob})

                                _, state_dis_fake, state_dis_real, state_dis_interpol, state = sess.run(
                                    [self.train_step_dis,
                                     self.next_state_dis_fake,
                                     self.next_state_dis_real,
                                     self.next_state_dis_interpol,
                                     self.next_state], feed_dict)
                # TEST:
                B_T, T_NUM_SEQ, _, _ = test_data.shape
                se = []
                predictions = []
                # Initial state:
                state = self._get_zero_state(B_T * self.args.n_mc_samples, self.args.gen_dim)
                for t_seq_ix in range(T_NUM_SEQ):
                    x_ = test_data[:, t_seq_ix]
                    mask = test_mask[:, t_seq_ix]
                    s_l_ = test_seq_len[:, t_seq_ix]

                    x_ = np.repeat(x_, self.args.n_mc_samples, axis=0)
                    s_l_ = np.repeat(s_l_, self.args.n_mc_samples, axis=0)
                    mask = np.repeat(mask, self.args.n_mc_samples, axis=0)

                    se_seq, state, prediction = sess.run(
                        [self.one_step_rmse, self.next_state, self.sample_dist], {
                            self.arrivals: x_,
                            self.state: state,
                            self.seq_len: s_l_,
                            self.mask: mask
                        })

                    prediction = np.concatenate(
                        (prediction, test_data[:, t_seq_ix]), axis=2)  # [B, T, MC + 1]
                    predictions.append(prediction)
                    se.append(se_seq)

                predictions = np.concatenate(predictions, axis=1)
                dist_mean = np.mean(
                    predictions[:, :, :self.args.n_mc_samples])

                rmse = np.mean(se[1:])
                service_times = np.asarray(predictions)
                ks_d, p = ks_2samp(
                    service_times[:, :, :-1].flatten(), service_times[:, :, -1].flatten())
                logger.info("RMSE on test set epoch {1}: {0:4f}, dist_mean: {2:4f}, d: {3:4f}, p: {4:4f}".format(
                    rmse, epoch, dist_mean, ks_d, p))
                if epoch % 10 == 0:
                    np.save(str(os.path.join(self.checkpoint_dir, "sample-iter-{}".format(epoch))), service_times)
                if ks_d < self.best_model_loss:
                    self.best_model_loss = ks_d
                    self._save_best_model(sess)

                self.store_training_epoch(sess, summary_writer, epoch,
                                          iteration, np.mean(mark_loss_set), (rmse, ks_d), np.mean(dis_loss_set))
        except KeyboardInterrupt:
            logger.info("Stop training at epoch {0}, iteration {1}".format(
                epoch, iteration))
        summary_writer.close()

    @staticmethod
    def _get_zero_state(batch_size, h_dim):
        c = np.zeros((batch_size, h_dim))
        h = np.zeros((batch_size, h_dim))

        return c, h


class AdversarialMempool(RnnHawkes):

    def __init__(self, args):

        super(AdversarialMempool, self).__init__(args)
        # self._build_placeholders()
        # self._build_recurrent()
        self._build_service_distribution()
        self.__build_training()
        # self._build_one_step_rmse()
        self.initialize = tf.global_variables_initializer()
        self.summary = tf.summary.merge_all()
        self._configure_logger(self.args.log.upper(),
                               "%(asctime)s - %(levelname)s - %(message)s")
        self.saver = tf.train.Saver(max_to_keep=100)
        self.sample_dist = self._get_predicted_marks()

    def _build_placeholders(self):

        super(AdversarialMempool, self)._build_placeholders()

        with tf.name_scope('Inputs'):
            self.arrivals_hidden = tf.placeholder(tf.float64, [None, None, self.args.arrivals_hidden_units])
            self.dropout_prob = tf.placeholder(
                tf.float64, shape=[], name="keep_prob_dropout")

    def _build_recurrent(self):
        cell_initializer = tf.initializers.random_uniform()
        if self.args.initializer == 'constant':
            cell_initializer = tf.initializers.constant(self.args.init_value)
        with tf.variable_scope(
                'PointProcess/RNN', initializer=cell_initializer):
            if self.args.cell == 'Basic':
                self._cell = tf.nn.rnn_cell.BasicRNNCell(self.args.hidden_units, tf.nn.relu)
            else:
                self._cell = tf.nn.rnn_cell.GRUCell(self.args.hidden_units, tf.nn.relu)
            self.reset_state = self._cell.zero_state(
                tf.shape(self.arrivals)[0], dtype=tf.float64)
            self.reset_state += tf.random_normal(
                tf.shape(self.reset_state), dtype=tf.float64)
            self.state = tf.placeholder(tf.float64,
                                        self.reset_state.get_shape(), "state")
            self.output, self.next_state = self._get_hidden_states(
                self.arrivals[:, :, :3], self.state, self.seq_len)

    def __build_training(self):

        with tf.name_scope("Train"):
            self.rmse_arrivals = tf.Variable(
                float("inf"), False, dtype=tf.float64, name="rmse_arrivals")
            self.rmse_marks = tf.Variable(
                float("inf"), False, dtype=tf.float64, name="rmse_marks")
            self.loss = tf.Variable(
                float("inf"), False, dtype=tf.float64, name="loss")
            self.loss_discriminator = tf.Variable(
                float("inf"), False, dtype=tf.float64, name="loss_dis")

            tf.summary.scalar(self.loss.op.name, self.loss)
            tf.summary.scalar(self.rmse_arrivals.op.name, self.rmse_arrivals)
            tf.summary.scalar(self.rmse_marks.op.name, self.rmse_marks)

            with tf.name_scope("Losses"):
                log_likelihood = self.log_likelihood()
                self.cost = tf.reduce_sum(
                    tf.multiply(log_likelihood, self.mask), name='cost')
                self.mark_loss = self._get_loss()
                self.dis_loss = self._get_discriminator_loss()

            self.__build_train_step()

    def __build_train_step(self):

        optimizer_1 = tf.train.AdamOptimizer(self.args.learning_rate)
        train_var_1 = tf.get_collection(
            tf.GraphKeys.GLOBAL_VARIABLES, scope="PointProcess")

        g_1 = tf.gradients(-self.cost, train_var_1)
        self.train_step = optimizer_1.apply_gradients(
            zip(g_1, train_var_1), global_step=self.global_step, name="train_step")

        optimizer_2 = tf.train.AdamOptimizer(
            learning_rate=self.args.learning_rate_gan,
            beta1=0.5,
            beta2=0.9)

        # Trainable variables within the GENERATOR

        train_gen_var = tf.get_collection(
            tf.GraphKeys.GLOBAL_VARIABLES, scope="GAN/generator")
        grad = tf.gradients(self.mark_loss, train_gen_var)
        self.train_step_gen = optimizer_2.apply_gradients(
            zip(grad, train_gen_var),
            name="train_step_gen")

        # Trainable variables within the DISCRIMINATOR

        train_dis_var = tf.get_collection(
            tf.GraphKeys.GLOBAL_VARIABLES, scope="GAN/discriminator")
        grad_dis = tf.gradients(self.dis_loss, train_dis_var)
        self.train_step_dis = optimizer_2.apply_gradients(
            zip(grad_dis, train_dis_var), name="train_step_dis")

    def _build_one_step_ahead_prediction(self):
        self.one_step_ahead_prediction_is = self._predict_one_step_ahead_inverse_sampling(
            self.arrivals[:, :, :3], self.state)

    def log_likelihood(self):
        with tf.name_scope('log_likelihood'):
            hidden = tf.concat((self.output, self.arrivals_hidden), -1)
            return self._arrivals_log_likelihood(self.arrivals, hidden)

    def _build_service_distribution(self):

        transaction_hidden_states = tf.concat((self.output, self.arrivals_hidden), axis=-1)
        covariates = self.arrivals[:, :, :3]
        b_real = tf.reshape(self.arrivals[:, :, -1],
                            [-1, 1])

        with tf.name_scope("build_generator-discriminator"):
            self._b_gen = self._generator(transaction_hidden_states, covariates)  # [B*T, 1], [B, T, H]

            self._logit_real = self._discriminator(transaction_hidden_states, covariates, b_real)  # [B*T, 1]

            self._logit_fake = self._discriminator(transaction_hidden_states, covariates, self._b_gen)  # [B*T, 1]

            self._grad_penalty = self._gradient_penalty_term(
                b_real, self._b_gen, transaction_hidden_states, covariates)

    def _generator(self, h, covariates):
        """
        GAN generator
        h (arrival hidden state) [B, T, H]
        covariates [B, T, D-1]
        """

        with tf.variable_scope("GAN/generator", reuse=tf.AUTO_REUSE):
            with tf.variable_scope("hidden_layers", reuse=tf.AUTO_REUSE):
                h = tf.reshape(h, [-1, self.args.arrivals_hidden_units + self.args.hidden_units])
                d_old = tf.reshape(covariates[:, :, 0], [-1, 1])
                delta_d = tf.reshape(covariates[:, :, 1], [-1, 1])
                b_old = tf.reshape(covariates[:, :, 2], [-1, 1])

                # embedding
                cell_initializer = tf.contrib.layers.xavier_initializer()
                h = tf.layers.dense(
                    h,
                    self.args.gen_dim,
                    kernel_initializer=cell_initializer,
                    name="layer_hidden_h")

                d_old = tf.layers.dense(
                    d_old,
                    self.args.gen_dim,
                    kernel_initializer=cell_initializer,
                    name="layer_hidden_d_old")

                delta_d = tf.layers.dense(
                    delta_d,
                    self.args.gen_dim,
                    kernel_initializer=cell_initializer,
                    name="layer_hidden_delta_d")

                b_old = tf.layers.dense(
                    b_old,
                    self.args.gen_dim,
                    kernel_initializer=cell_initializer,
                    name="layer_hidden_b_old")

                # sampling prior normal distr.
                z = tf.random_normal(tf.shape(h), 0.0, 1.0, dtype=tf.float64)
                z = tf.layers.dense(
                    z,
                    self.args.gen_dim,
                    use_bias=False,
                    kernel_initializer=cell_initializer,
                    name="layer_noise_0")

                # hidden state
                h = tf.nn.leaky_relu(h + d_old + delta_d + b_old + z)  # [B*T, D_GAN]

                if self.args.n_layers_gen > 1:
                    for i in range(self.args.n_layers_gen - 1):
                        h = tf.layers.dense(
                            h,
                            self.args.gen_dim,
                            kernel_initializer=cell_initializer,
                            name="layer_hidden_{}".format(i + 1))
                        if self.args.noisy_layers:
                            z = tf.random_normal(
                                tf.shape(h), 0.0, 1.0, dtype=tf.float64)
                            z = tf.layers.dense(
                                z,
                                self.args.gen_dim,
                                use_bias=False,
                                kernel_initializer=cell_initializer,
                                name="layer_noise_{}".format(i + 1))
                            h = h + z
                        if self.args.dropout:
                            h = tf.nn.dropout(h, self.dropout_prob)
                        h = tf.nn.leaky_relu(h)  # [B*T, D_GAN]

            with tf.variable_scope("generator_output", reuse=tf.AUTO_REUSE):
                h = tf.layers.dense(h, 1, kernel_initializer=cell_initializer)

        return tf.exp(h)  # [B*T, 1]

    def _discriminator(self, h, covariates, b_new):
        """
        GAN discriminator
        h (arrival hidden state) [B, T, H]
        covariates [B, T, D-3]
        s (service time) [B*T, 1]
        """
        with tf.variable_scope("GAN/discriminator", reuse=tf.AUTO_REUSE):
            with tf.variable_scope("hidden_layers", reuse=tf.AUTO_REUSE):
                h = tf.reshape(h, [-1, self.args.arrivals_hidden_units + self.args.hidden_units])
                d_old = tf.reshape(covariates[:, :, 0], [-1, 1])
                delta_d = tf.reshape(covariates[:, :, 1], [-1, 1])
                b_old = tf.reshape(covariates[:, :, 2], [-1, 1])

                # embedding
                cell_initializer = tf.contrib.layers.xavier_initializer()
                h = tf.layers.dense(
                    h,
                    self.args.dis_dim,
                    kernel_initializer=cell_initializer,
                    name="layer_hidden_h")

                d_old = tf.layers.dense(
                    d_old,
                    self.args.dis_dim,
                    kernel_initializer=cell_initializer,
                    name="layer_hidden_d_old")

                delta_d = tf.layers.dense(
                    delta_d,
                    self.args.dis_dim,
                    kernel_initializer=cell_initializer,
                    name="layer_hidden_delta_d")

                b_old = tf.layers.dense(
                    b_old,
                    self.args.dis_dim,
                    kernel_initializer=cell_initializer,
                    name="layer_hidden_b_old")

                b_new = tf.layers.dense(
                    b_new,
                    self.args.dis_dim,
                    use_bias=False,
                    kernel_initializer=cell_initializer,
                    name="layer_hidden_b_new")

                # hidden state
                h = tf.nn.leaky_relu(h + d_old + delta_d + b_old + b_new)  # [B*T, D_GAN]

                if self.args.n_layers_dis > 1:
                    for i in range(self.args.n_layers_dis - 1):
                        h = tf.layers.dense(
                            h,
                            self.args.dis_dim,
                            kernel_initializer=cell_initializer,
                            name="layer_hidden_{}".format(i + 1))
                        if self.args.dropout:
                            h = tf.nn.dropout(h, self.dropout_prob)
                        h = tf.nn.leaky_relu(h)  # [B*T, D_GAN]

            with tf.variable_scope(
                    "discriminator_output", reuse=tf.AUTO_REUSE):
                h = tf.layers.dense(h, 1, kernel_initializer=cell_initializer)

        return h  # [B*T, 1]

    def _gradient_penalty_term(self, s_real, s_gen, h, covariates):

        with tf.name_scope("gradient_penalty_term"):
            eta = tf.random_uniform(
                tf.shape(s_real), minval=0.0, maxval=0.1, dtype=tf.float64)

            s_interpol = eta * s_real + (1 - eta) * s_gen
            logit_interpol = self._discriminator(h, covariates, s_interpol)

            grad_penalty = tf.gradients(logit_interpol, s_interpol)[
                0]  # tf.gradients returns a list

            return tf.abs(grad_penalty)

    def _get_discriminator_loss(self):
        with tf.name_scope("adversarial_loss"):
            with tf.name_scope("discriminator_component"):
                shape = tf.shape(self.arrivals)
                batch_size = shape[0]
                seq_length = shape[1]

                dis_loss = -(self._logit_real - self._logit_fake)

                with tf.name_scope("gradient_penalty"):
                    dis_loss = dis_loss + self.args.lambda_grad_penalty * \
                               tf.square(tf.maximum(tf.cast(0.0, tf.float64),
                                                    self._grad_penalty - 1.0))

                dis_loss = tf.reshape(dis_loss, [batch_size, seq_length])
                dis_loss = tf.multiply(dis_loss, self.mask)
                dis_loss = tf.reduce_sum(dis_loss)

            return dis_loss

    def _get_loss(self):
        with tf.name_scope("adversarial_loss"):
            with tf.name_scope("generator_component"):
                shape = tf.shape(self.arrivals)
                batch_size = shape[0]
                seq_length = shape[1]

                gen_loss_1 = -self._logit_fake

                gen_loss_1 = tf.reshape(gen_loss_1, [batch_size, seq_length])
                gen_loss_1 = tf.multiply(gen_loss_1, self.mask)

            return tf.reduce_sum(gen_loss_1)

    def _predict_mark(self, covariates):
        """
        Args:
            output:
            covariates:
        """

        h = tf.concat((self.output, self.arrivals_hidden), axis=-1)
        shape = tf.shape(self.arrivals_hidden)
        next_mark = self._generator(h, covariates)
        next_mark = tf.reshape(next_mark, (shape[0], shape[1]))

        return next_mark

    def _predict_one_step_ahead_inverse_sampling(self, x, last_state):
        """

        Args:
            x:
            last_state:

        Returns:

        """
        with tf.name_scope('inverse_sampling'):
            states, _ = self._get_hidden_states(x, last_state, None)
            x_shape = tf.shape(x)
            y = tf.random_uniform([
                x_shape[0], x_shape[1],
                tf.cast(self.args.inverse_sampling_size, tf.int32)
            ])
            states = tf.concat((states, self.arrivals_hidden), -1)
            alpha = self._past_influence(states) + self.b_t

            alpha = tf.expand_dims(alpha, 2)
            log_inv_sample = tf.cast(
                tf.log(tf.divide(-1., y - 1.)), tf.float64)
            inv_sample = tf.divide(
                tf.log(
                    tf.abs(
                        tf.multiply(
                            self.w_t, log_inv_sample + tf.divide(
                                tf.exp(alpha), self.w_t)))) - alpha, self.w_t)
            return x[:, :, 0] + tf.reduce_mean(inv_sample, axis=2)

    def _get_predicted_marks(self):

        h = tf.concat((self.output, self.arrivals_hidden), axis=-1)
        covariates = self.arrivals[:, :, :-1]
        shape = tf.shape(self.arrivals)

        next_mark = self._generator(h, covariates)

        next_mark = tf.reshape(next_mark, (shape[0], shape[1]))

        next_mark = tf.multiply(self.mask, next_mark)

        next_mark = tf.split(next_mark, self.args.batch_size_test, axis=0)
        next_mark = tf.stack(next_mark)  # [B, MC, T]
        next_mark = tf.transpose(next_mark, [0, 2, 1])

        return next_mark

    def _build_one_step_rmse(self):

        # Arrivals error:
        y_predicted = self.one_step_ahead_prediction_is
        y_actual = self.arrivals[:, :, -2]

        y_predicted = tf.multiply(y_predicted, self.mask)
        se_arrival = tf.abs(tf.subtract(y_predicted, y_actual))

        marks_real = self.arrivals[:, :, -1]
        covariates = self.arrivals[:, :, :-1]

        marks_predicted = self._predict_mark(covariates)
        marks_predicted = tf.multiply(self.mask, marks_predicted)

        # Marks error:
        # Averaged values:
        # 1.- predictions:
        marks_predicted = tf.split(marks_predicted,
                                   self.args.batch_size_test,
                                   axis=0)  # [B, MC, T]
        marks_predicted = tf.stack(marks_predicted)
        marks_predicted = tf.transpose(marks_predicted, [0, 2, 1])  # [B, T, MC]
        marks_predicted = tf.reduce_mean(marks_predicted, axis=2)
        # 2.- real:
        marks_real = tf.split(marks_real,
                              self.args.batch_size_test,
                              axis=0)  # [B, MC, T]
        marks_real = tf.stack(marks_real)
        marks_real = tf.transpose(marks_real, [0, 2, 1])  # [B, T, MC]
        marks_real = tf.reduce_mean(marks_real, axis=2)

        se_mark = tf.abs(marks_predicted - marks_real)

        self.one_step_rmse = tf.tuple((se_arrival, se_mark))

    def intensity(self):
        with tf.name_scope('intensity'):
            hidden_states = tf.concat((self.output, self.arrivals_hidden), axis=-1)
            return self._intensity(self.arrivals, hidden_states)

    def __log_intensity(self, x, hidden_states):
        """
        Args:
            x: [B, T, D]
            hidden_states: [B, T, H]

        Returns:
            [B, T]
        """
        with tf.name_scope('log_intensity'):
            past_influence = self._past_influence(hidden_states)
            current_influence = tf.multiply(self.w_t, x[:, :, -1] - x[:, :, 0],
                                            'current_influence')
            return past_influence + current_influence + self.b_t

    def train(self, train_set, test_set, hidden_train, hidden_test, sess):
        """

        Args:
            train_data:
            test_data:
            seq_len:
            mask:
            sess:

        Returns:

        """
        train_data, train_seq_len, train_mask = train_set
        test_data, test_seq_len, test_mask = test_set
        sess.run(self.initialize)
        summary_writer = self.summary_writer(self.tensorboard_dir, sess)
        NUM_BATCHES, _, NUM_SEQ, _, _ = train_data.shape
        TOTAL_ITERATIONS = self.args.num_epochs * NUM_BATCHES * (NUM_SEQ - 1.)
        num_epochs = self.args.num_epochs
        try:
            for epoch in range(num_epochs):

                print "SHAPE train DATA", train_data[0, :, 0].shape

                state = sess.run(self.reset_state, {self.arrivals: train_data[0, :, 0]})

                print "PRINTING", train_data[0, :, 0].shape
                print "PRINTING", state.shape
                print "PRINTING", type(state), type(state[0])

                for batch_ix in range(NUM_BATCHES):
                    costs = []
                    for seq_ix in range(NUM_SEQ):
                        x = train_data[batch_ix, :, seq_ix]
                        s_l = train_seq_len[batch_ix, :, seq_ix]
                        m = train_mask[batch_ix, :, seq_ix]
                        h = hidden_train[batch_ix, :, seq_ix]
                        feed_dict = {
                            self.arrivals: x,
                            self.state: state,
                            self.seq_len: s_l,
                            self.mask: m,
                            self.arrivals_hidden: h,
                            self.dropout_prob: self.args.dropout_prob
                        }

                        for _ in range(self.args.n_discriminator_updates):
                            dis_loss, _ = sess.run(
                                [self.dis_loss, self.train_step_dis], feed_dict)

                        for _ in range(self.args.n_generator_updates):
                            mark_loss, _ = sess.run([self.mark_loss,
                                                     self.train_step_gen],
                                                    feed_dict)

                        fetches = [
                            self.cost, self.next_state, self.global_step
                        ]

                        if seq_ix > 0:
                            fetches.append(self.train_step)
                            cost, state, iteration, _ = sess.run(
                                fetches, feed_dict)
                        else:
                            cost, state, iteration = sess.run(
                                fetches, feed_dict)

                        if self.logging_interval(iteration, self.args.logging_interval):
                            logger.info(
                                "Epoch {0} ({1:.2%}) Step {2}: (gen) {3:4f}, (crit) {4:4f}, (rpp) {5:4f}"
                                    .format(epoch,
                                            iteration / TOTAL_ITERATIONS,
                                            iteration, mark_loss, dis_loss, cost))

                        costs.append(cost)

                # TEST
                N, T_NUM_SEQ, _, _ = test_data.shape
                se_arr = []
                se_marks = []
                predictions_marks = []
                state = sess.run(self.reset_state, {
                    self.arrivals: np.repeat(test_data[:, 0], self.args.n_mc_samples, axis=0)})

                print "PRINTING", test_data[:, 0].shape
                print "PRINTING", state.shape
                print "PRINTING", type(state), type(state[0])

                for t_seq_ix in range(T_NUM_SEQ):
                    x_ = test_data[:, t_seq_ix]
                    h_ = hidden_test[:, t_seq_ix]
                    s_l_ = test_seq_len[:, t_seq_ix]
                    mask = test_mask[:, t_seq_ix]

                    x_ = np.repeat(x_, self.args.n_mc_samples, axis=0)
                    h_ = np.repeat(h_, self.args.n_mc_samples, axis=0)
                    s_l_ = np.repeat(s_l_, self.args.n_mc_samples, axis=0)
                    mask = np.repeat(mask, self.args.n_mc_samples, axis=0)

                    se_seq, state, prediction = sess.run(
                        [self.one_step_rmse, self.next_state, self.sample_dist], {
                            self.arrivals: x_,
                            self.state: state,
                            self.seq_len: s_l_,
                            self.mask: mask,
                            self.arrivals_hidden: h_,
                            self.dropout_prob: 1.0
                        })

                    prediction = np.concatenate(
                        (prediction, test_data[:, t_seq_ix]), axis=2)  # [B, T, MC + 1]
                    predictions_marks.append(prediction)
                    se_arr.append(se_seq[0])
                    se_marks.append(se_seq[1])
                predictions_marks = np.concatenate(predictions_marks, axis=1)
                dist_mean = np.mean(
                    predictions_marks[:, :, :self.args.n_mc_samples])
                rmse_arr = np.mean(se_arr[1:])
                rmse_mark = np.mean(se_marks[1:])
                predictions_marks = np.asarray(predictions_marks)
                ks_d, p = ks_2samp(
                    predictions_marks[:, :, :-1].flatten(), predictions_marks[:, :, -1].flatten())
                logger.info(
                    "RMSE on test set epoch {1}: (arri) {0:4f}, (mark) {5:4f}, dist_mean: {2:4f}, d: {3:4f}, p: {4:4f}".format(
                        rmse_arr, epoch, dist_mean, ks_d, p, rmse_mark))

                if epoch % 10 == 0:
                    np.save(str(os.path.join(self.checkpoint_dir, "sample-iter-{}".format(epoch))), predictions_marks)
                if ks_d < self.best_model_loss:
                    self.best_model_loss = ks_d
                    self._save_best_model(sess)

                self.store_training_epoch(sess, summary_writer, epoch,
                                          iteration, np.sum(costs), (rmse_arr, ks_d))

        except KeyboardInterrupt:
            logger.info("Stop training at epoch {0}, iteration {1}".format(
                epoch, iteration))
        summary_writer.close()
