'''
Created on Oct 1, 2018

@author: anon
'''
import copy
import datetime
import json
import logging
import math
import os
from abc import ABCMeta, abstractmethod

import numpy as np
import tensorflow as tf

from rnpp.npp import logger
from rnpp.utils.etc.quadratures import quadratures
from rnpp.utils.etc.utils import Bunch

# tf.compat.v1.nn.rnn_cell
# _Linear = core_rnn_cell._Linear

def generate_initialized_variable(shape,
                                  init_value=1.,
                                  name=None,
                                  dtype=tf.float64,
                                  constant_initializer=False):
    """
    Args:
        shape:
        init_value:
        name:
        dtype:
        constant_initializer:
    """
    if constant_initializer:
        initialization = tf.compat.v1.ones(shape, dtype) * init_value
    else:
        initialization = tf.compat.v1.random_uniform(shape, -init_value, init_value,
                                           dtype)
    return tf.compat.v1.Variable(initialization, name=name)


def generate_initialized_scalar_variable(name,
                                         init_value,
                                         dtype=tf.float64,
                                         constant_initializer=False):
    """
    Args:
        name:
        init_value:
        dtype:
        constant_initializer:
    """
    if constant_initializer:
        initialization = tf.compat.v1.ones((), dtype) * init_value
    else:
        initialization = tf.compat.v1.random_uniform((),
                                           -init_value,
                                           init_value,
                                           dtype=dtype)
    return tf.compat.v1.Variable(initialization, name=name)


class Model:
    __metaclass__ = ABCMeta

    def __init__(self, args):
        """
        Args:
            args:
        """
        if 'restore' not in args:
            args['restore'] = False
        self.args = copy.deepcopy(Bunch(args))
        self.loss = None
        self.loss_discriminator = None
        self.rmse_arrivals = None
        self.rmse_marks = None
        self._one = tf.compat.v1.constant(1, dtype=tf.int32)
        self.summary = None
        if self.args.numpy_seed is not None:
            np.random.seed(self.args.numpy_seed)
        if self.args.tf_seed is not None:
            tf.random.set_seed(self.args.tf_seed)
            tf.compat.v1.set_random_seed(self.args.tf_seed)
        self.global_step = tf.compat.v1.Variable(0, name='global_step', trainable=False)
        self.global_step_gan = tf.compat.v1.Variable(
            0, name='global_step_gan', trainable=False)
        self.best_model_loss = float("inf")
        self._build_placeholders()
        self._prepare_dirs()
        self.saver = None

    @abstractmethod
    def _build_recurrent(self):
        return NotImplemented

    @abstractmethod
    def _build_placeholders(self):
        return NotImplemented

    @abstractmethod
    def _build_one_step_ahead_prediction(self):
        return NotImplemented

    @abstractmethod
    def _predict_one_step_ahead(self, x, last_state):
        """
        Args:
            x:
            last_state:
        """
        return NotImplemented

    @abstractmethod
    def _build_one_step_rmse(self):
        return NotImplemented

    @abstractmethod
    def _arrivals_log_likelihood(self, x, hidden_state):
        """
        Args:
            x:
            hidden_state:
        """
        return NotImplemented

    @abstractmethod
    def train(self, train_set, test_set, sess):
        """
        Args:
            train_data:
            test_data:
            sess:
        """
        return NotImplemented

    @abstractmethod
    def inference(self, x, sess):
        """
        Args:
            x:
            sess:
        """
        return NotImplemented

    def sequences_length(self, sequences):
        """
        Returns the length of each sequence. The length of a sequence is equal to the number of non-zero elements.
        :param sequences: expected [BxMAX_TxD] shape
        :return:
        """
        non_zero = tf.sign(tf.reduce_max(tf.abs(sequences), 2))
        sequences_length = tf.reduce_sum(non_zero, 1)
        sequences_length = tf.cast(sequences_length, tf.int32)
        return sequences_length

    def _configure_logger(self, level, logging_format):
        logger.setLevel(level)
        sh = logging.StreamHandler()
        fh = logging.FileHandler(os.path.join(self.tensorboard_dir, "raw.log"))
        sh.setLevel(level)
        fh.setLevel(level)
        formatter = logging.Formatter(logging_format)
        sh.setFormatter(formatter)
        fh.setFormatter(formatter)
        logger.addHandler(sh)
        logger.addHandler(fh)

    def _prepare_dirs(self):
        if self.args.restore:
            start_time = self.args.restore_dir.split(os.sep)[-1]
        else:
            start_time = datetime.datetime.now().strftime("%d%m_%H%M%S%f")
        self.checkpoint_dir = os.path.join(self.args.store_dir, self.args.name,
                                           start_time)
        self.tensorboard_dir = os.path.join(self.args.log_dir, self.args.name,
                                            start_time)
        if self.args.restore is False:
            os.makedirs(self.checkpoint_dir)
            os.makedirs(self.tensorboard_dir)

    def _save_best_model(self, session):
        file_name = os.path.join(self.checkpoint_dir, "best-model")
        logger.info("Saving best model ...")
        tf.compat.v1.train.Saver().save(session, file_name)

    def store_training_epoch(self,
                             session,
                             summary_writer,
                             epoch,
                             iteration,
                             loss,
                             rmse,
                             loss_dis=None):
        """
        Args:
            session:
            summary_writer:
            iteration:
            loss:
            rmse:
        """
        if loss_dis is not None:
            session.run([
                self.loss.assign(loss),
                self.loss_discriminator.assign(loss_dis),
                self.rmse_arrivals.assign(rmse[0]),
                self.rmse_marks.assign(rmse[1])
            ])
        else:
            session.run([
                self.loss.assign(loss),
                self.rmse_arrivals.assign(rmse[0]),
                self.rmse_marks.assign(rmse[1])
            ])
        summary_writer.add_summary(
            session.run(self.summary), global_step=iteration)
        summary_writer.flush()
        if epoch % 10 == 0:
            model_filename = os.path.join(self.checkpoint_dir,
                                          "model-iteration")
            self.saver.save(session, model_filename, epoch)
            self._write_model_parameters(self.checkpoint_dir)
            logger.info("Saving checkpoint: {} ...".format(model_filename))

    def _write_model_parameters(self, model_directory):
        """
        Args:
            model_directory:
        """
        with open(self._parameters_file(model_directory), "w") as f:
            json.dump(vars(self.args), f, indent=4)

    @classmethod
    def restore(cls, session, path):
        """Restore a previously trained model

        Args:
            session (TensorFlow Session): session into which to restore the
                model
            model_directory (str): directory to which the model was saved

        Returns:
            RNN: trained model
        """
        model_directory,_ = os.path.split(path)

        with open(cls._parameters_file(model_directory)) as f:
            parameters = json.load(f)
        parameters['restore'] = True
        parameters['restore_dir'] = model_directory
        model = cls(parameters)

        tf.compat.v1.train.Saver().restore(session, path)
        logger.info("Loading model " + path)

        return model

    @staticmethod
    def logging_interval(iteration, interval):
        """
        Args:
            iteration:
            interval:
        """
        return interval is not None and iteration > 1 and iteration % interval == 0

    @staticmethod
    def summary_writer(summary_directory, session):
        """
        Args:
            summary_directory:
            session:
        """

        class NullSummaryWriter(object):

            def add_summary(self, *args, **kwargs):
                pass

            def flush(self):
                pass

            def close(self):
                pass

        if summary_directory is not None:
            return tf.compat.v1.summary.FileWriter(summary_directory, session.graph)
        else:
            return NullSummaryWriter()

    @staticmethod
    def _parameters_file(model_directory):
        """
        Args:
            model_directory:
        """
        return os.path.join(model_directory, "parameters.json")


class RnnHawkes(Model):

    def __init__(self, args):
        """
        Args:
            args:
        """
        super(RnnHawkes, self).__init__(args)
        self._build_recurrent()
        self.__build_intensity_variables()
        self._build_event_type_generation_logits()
        self.__build_training()
        self._build_one_step_ahead_prediction()
        self._build_one_step_rmse()
        self.initialize = tf.compat.v1.global_variables_initializer()
        self.summary = tf.compat.v1.summary.merge_all()
        self._configure_logger(self.args.log.upper(),
                               "%(asctime)s - %(levelname)s - %(message)s")
        self.saver = tf.compat.v1.train.Saver(max_to_keep=100)

    def _build_placeholders(self):
        """The shape of the input data is [B, T, D] B - is batch size T - is the
        window size of BPTT D - is the dimension of the data in this case
        D=2(ex. event time and inter arrival time) :return:
        """
        with tf.compat.v1.name_scope('Inputs'):
            self.arrivals = tf.compat.v1.placeholder(
                tf.compat.v1.float64, [None, None, self.args.input_dim], 'input_data')
            self.seq_len = tf.compat.v1.placeholder(tf.compat.v1.int32, [None], 'seq_len')
            self.mask = tf.compat.v1.placeholder(tf.compat.v1.float64, [None, None], 'mask')

    def _build_recurrent(self):
        cell_initializer = tf.compat.v1.initializers.random_uniform()
        if self.args.initializer == 'constant':
            cell_initializer = tf.compat.v1.initializers.constant(self.args.init_value)

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

    def _build_event_type_generation_logits(self):
        if self.args.M > 0:
            self.event_logits = self._event_type_generation_logits(self.output)

    def _event_type_generation_logits(self, hidden_state):
        # hidden_state: [B, T, h]
        # V_y: [h, M, K]
        logits = tf.compat.v1.tensordot(hidden_state, self.V_y, [[2], [0]],
                              'events_logits')
        logits += self.b_y
        return logits  # [B, T, M, K]

    def _get_hidden_states(self, x, initial_state, seq_len):
        """
        Args:
            x:
            initial_state:
        """
        output, next_state = tf.compat.v1.nn.dynamic_rnn(self._cell, x, seq_len,
                                               initial_state)
        output = self.__put_last_state_to_output(output, initial_state)
        return output, next_state

    def __put_last_state_to_output(self, hidden_states, last_state):
        """
        Args:
            hidden_states:
            last_state:
        """
        h_j = last_state[:, None]
        hidden_states = hidden_states[:, :-1]
        hidden_states = tf.compat.v1.concat([h_j, hidden_states], axis=self._one)
        return hidden_states

    def __build_intensity_variables(self):
        with tf.compat.v1.name_scope('PointProcess/intensity_parameters'):
            h_units = self.args.hidden_units + self.args.__dict__.get('arrivals_hidden_units', 0)
            is_constant_initializer = self.args.initializer == 'constant'
            self.v_t = generate_initialized_variable(
                (h_units, 1), self.args.init_value, 'V_t',
                tf.compat.v1.float64, is_constant_initializer)
            self.w_t = generate_initialized_scalar_variable(
                'w_t', self.args.init_value, tf.compat.v1.float64,
                is_constant_initializer)
            self.b_t = generate_initialized_scalar_variable(
                'b_t', self.args.init_value, tf.compat.v1.float64,
                is_constant_initializer)

            # Train static mark model
            if self.args.M > 0:
                self.V_y = generate_initialized_variable(
                    (h_units, self.args.M, self.args.K),
                    self.args.init_value, 'V_y', tf.compat.v1.float64,
                    is_constant_initializer)
                self.b_y = generate_initialized_variable(
                    (
                        self.args.M,
                        self.args.K,
                    ), self.args.init_value, 'b_y', tf.float64,
                    is_constant_initializer)

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

            log_likelihood = self.log_likelihood()
            self.cost = tf.reduce_sum(
                tf.multiply(log_likelihood, self.mask), name='cost')
            tf.compat.v1.summary.scalar(self.loss.op.name, self.loss)
            tf.compat.v1.summary.scalar(self.rmse_arrivals.op.name, self.rmse_arrivals)
            tf.compat.v1.summary.scalar(self.rmse_marks.op.name, self.rmse_marks)
            self.__build_train_step()

    def __build_train_step(self):

        optimizer = tf.compat.v1.train.AdamOptimizer(self.args.learning_rate)

        # train_var = tf.trainable_variables()

        # All trainable variables associated to the RPP

        train_var = tf.compat.v1.get_collection(
            tf.compat.v1.GraphKeys.GLOBAL_VARIABLES, scope="PointProcess")

        g = tf.compat.v1.gradients(-self.cost, train_var)
        self.train_step = optimizer.apply_gradients(
            zip(g, train_var), global_step=self.global_step, name="train_step")

    def intensity(self):
        with tf.compat.v1.name_scope('intensity'):
            return self._intensity(self.arrivals, self.output)

    def _intensity(self, x, hidden_state):
        """
        Args:
            x:
            hidden_state:
        """
        return tf.exp(self.__log_intensity(x, hidden_state))

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

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

    def _past_influence(self, hidden_states):
        """v_t: [H, 1] :param hidden_states: [B, T, H] :return: [B, T]

        Args:
            hidden_states:
        """
        return tf.compat.v1.squeeze(
            tf.compat.v1.tensordot(hidden_states, self.v_t, [[2], [0]],
                         'past_influence'), 2)

    def log_likelihood(self):
        with tf.compat.v1.name_scope('log_likelihood'):
            if self.args.M > 0:
                return self._arrivals_log_likelihood(self.arrivals, self.output) + \
                       self._marks_log_likelihood(
                           self.arrivals[:, :, 2:], self.event_logits)
            else:
                return self._arrivals_log_likelihood(self.arrivals,
                                                     self.output)

    def _likelihood(self, x, hidden_state):
        """
        Args:
            x:
            hidden_state:
        """
        return tf.compat.v1.exp(self._arrivals_log_likelihood(x, hidden_state))

    def _arrivals_log_likelihood(self, x, hidden_state):
        """ Calculates the log likelihood according to equation 14.

        Args:
            x:
            hidden_state:
        """
        log_intensity = self.__log_intensity(x, hidden_state)
        past_influence = self._past_influence(hidden_state)
        return log_intensity + tf.divide(
            tf.compat.v1.exp(past_influence + self.b_t), self.w_t) - tf.compat.v1.divide(
            tf.compat.v1.exp(log_intensity), self.w_t)

    def _marks_log_likelihood(self, x, logits):
        """
        :param x: [B, T, M]
        :param logits: [B, T, M, K]
        :return: [B, T]
        """
        log_p = tf.compat.v1.log(tf.compat.v1.nn.softmax(logits, axis=-1))  # [B, T, M, K]
        indices = tf.compat.v1.cast(tf.expand_dims(x, axis=-1), tf.int32)  # [B, T, M, 1]
        log_p = tf.compat.v1.squeeze(
            tf.compat.v1.batch_gather(log_p, indices), axis=-1)  # [B, T, M]
        log_p = tf.compat.v1.reduce_sum(log_p, axis=-1)
        return log_p

    def _build_one_step_rmse(self):
        y_predicted = self.one_step_ahead_prediction_is
        y_actual = self.arrivals[:, 1:, 0]

        y_predicted = tf.compat.v1.multiply(y_predicted, self.mask[:, 1:])
        se = tf.compat.v1.square(tf.compat.v1.subtract(y_predicted, y_actual))

        self.one_step_rmse = se

    def _build_one_step_ahead_prediction(self):
        self.one_step_ahead_prediction = self._predict_one_step_ahead(
            self.arrivals, self.state)
        self.one_step_ahead_prediction_is = self._predict_one_step_ahead_inverse_sampling(
            self.arrivals, self.state)

    def _predict_one_step_ahead(self, x, last_state):
        """
        Args:
            x:
            last_state:
        """
        num_q_p = self.args.num_quadrature_nodes
        states, _ = self._get_hidden_states(x, last_state, self.seq_len)
        states = tf.compat.v1.tile(states[:, :, None], [1, 1, num_q_p, 1])
        states = tf.compat.v1.reshape(states, (-1, num_q_p, self.args.hidden_units))

        def integrand(q_p):
            q_p_delta = tf.compat.v1.concat((x[:, :, :1], q_p), 2)
            q_p_delta = q_p_delta[:, :, 1:] - q_p_delta[:, :, :-1]
            q_p = tf.compat.v1.concat((q_p[:, :, :, None], q_p_delta[:, :, :, None]), 3)

            T = tf.compat.v1.shape(q_p)[1]
            q_p = tf.compat.v1.reshape(q_p, (-1, num_q_p, self.args.input_dim))
            likelihoods = tf.compat.v1.multiply(q_p[:, :, 0],
                                      self._likelihood(q_p, states))
            likelihoods = tf.compat.v1.reshape(likelihoods, (-1, T, num_q_p))
            return likelihoods

        a = x[:, :, 0]
        b = a + self.args.quadrature_window
        predictions = quadratures(integrand, a, b)
        return predictions

    def predict(self, x, sess, num_points=10):
        """
        Args:
            x:
            sess:
            num_points:
        """
        B = x.shape[0]
        L = x.shape[1]
        seq_len = np.ones(B) * L
        predictions = []
        state = sess.run(self.reset_state)
        feed_dict = {
            self.arrivals: x,
            self.state: state,
            self.seq_len: seq_len
        }

        next_state = sess.run(self.next_state, feed_dict)
        t = x[:, -1:]
        for _ in range(num_points):
            t_j = sess.run(self.__predict_next_point(t, next_state))
            next_state = sess.run(
                self.next_state, {
                    self.arrivals: t_j,
                    self.state: next_state,
                    self.seq_len: np.ones(B)
                })
            t = t_j
            predictions.append(t_j[:, :, 0])
        return np.concatenate(predictions, 1)

    def __predict_next_point(self, t, next_state):
        """
        Args:
            t:
            next_state:
        """
        t_j = self._predict_one_step_ahead(t, next_state)
        d_t = t_j - t[:, :, 0]
        t_j = tf.compat.v1.concat((t_j, d_t), 1)
        return t_j[:, None]

    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.compat.v1.shape(x)
            y = tf.compat.v1.random_uniform([
                x_shape[0], x_shape[1],
                tf.compat.v1.cast(self.args.inverse_sampling_size, tf.int32)
            ])
            alpha = self._past_influence(states) + self.b_t

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

    def inference(self, x, sess):
        """
        :param x tensor with size [N, L, D] :param sess

        N is the number of sequences L is the length of a sequence D is the
        dimension of each arrival

        Args:
            x:
            sess:
        """
        seq_len = np.ones(x.shape[0]) * x.shape[1]
        state = sess.run(self.reset_state)
        intensity = sess.run(self.intensity(), {
            self.arrivals: x,
            self.state: state,
            self.seq_len: seq_len
        })
        return intensity

    def extract_hidden_states(self, train_set, test_set, sess):
        """
        Obtain the hidden states of the arrival network equation 16.

        :param data tensor with size [N, L, D] -  N is the number of sequences L is the length of a sequence D is the dimension of each arrival
        :param sess
        """
        train_data, train_seq_len, _ = train_set
        test_data, test_seq_len, test_mask = test_set
        NUM_BATCHES, _, NUM_SEQ, _, _ = train_data.shape

        test_states = []
        train_states = []
        state = sess.run(self.reset_state, {
            self.arrivals: train_data[0, :, 0]})
        for batch_ix in range(NUM_BATCHES):
            states = []
            for seq_ix in range(NUM_SEQ):
                x = train_data[batch_ix, :, seq_ix]
                s_l = train_seq_len[batch_ix, :, seq_ix]

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

        N, T_NUM_SEQ, _, _ = test_data.shape
        state = sess.run(self.reset_state, {
            self.arrivals: test_data[:, 0]})
        state += state[:N]
        for seq_ix in range(T_NUM_SEQ):
            output, state = sess.run(
                [self.output, self.next_state], {
                    self.arrivals: test_data[:, seq_ix],
                    self.state: state,
                    self.seq_len: test_seq_len[:, seq_ix],
                    self.mask: test_mask[:, seq_ix]
                })
            test_states.append(np.expand_dims(output, 1))

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

        return train_states, test_states

    def extract_hidden_states_test(self, test_set, sess):
        """
        Obtain the hidden states of the arrival network equation 16.

        :param data tensor with size [N, L, D] -  N is the number of sequences L is the length of a sequence D is the dimension of each arrival
        :param sess
        """
        test_data, test_seq_len, test_mask = test_set

        test_states = []

        N, T_NUM_SEQ, _, _ = test_data.shape
        state = sess.run(self.reset_state, {
            self.arrivals: test_data[:, 0]})
        state += state[:N]
        for seq_ix in range(T_NUM_SEQ):
            output, state = sess.run(
                [self.output, self.next_state], {
                    self.arrivals: test_data[:, seq_ix],
                    self.state: state,
                    self.seq_len: test_seq_len[:, seq_ix],
                    self.mask: test_mask[:, seq_ix]
                })
            test_states.append(np.expand_dims(output, 1))

        test_states = np.concatenate(test_states, 1)

        return test_states

    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]})
                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]
                        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 = []
                state = sess.run(self.reset_state, {
                    self.arrivals: test_data[:, 0]})
                state += state[:N]
                fetches = [self.one_step_rmse, self.next_state]
                for t_seq_ix in range(T_NUM_SEQ):
                    se_seq, state = 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]
                        })
                    se.append(se_seq)
                rmse = np.sqrt(np.mean(se[1:]))
                logger.info(
                    "RMSE on test set epoch {1}: arrival {0:4f}, marks {2:4f}".
                        format(rmse, epoch, 0))
                if math.isnan(rmse):
                    raise KeyboardInterrupt
                self.store_training_epoch(sess, summary_writer, epoch,
                                          iteration, np.sum(costs), (rmse, 0))

                if rmse < self.best_model_loss:
                    self.best_model_loss = rmse
                    self._save_best_model(sess)
        except KeyboardInterrupt:
            logger.info("Stop training at epoch {0}, iteration {1}".format(
                epoch, iteration))
        summary_writer.close()
