'''
Created on Oct 1, 2018

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

import numpy as np
import tensorflow as tf
from scipy.stats import ks_2samp

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

def update_system_index(system_index_old):
    """
    obtain the next event by comparing all next arrivals and departures

    load new files if index are exceeded, and updates system index

    :return:
    """
    system_index = copy.deepcopy(system_index_old)

    data_dir = system_index["data_dir"]
    current_arrival_file_index = system_index["current_arrival_file_index"]
    arrivals_index_in_file = system_index["arrivals_index_in_file"]
    current_departure_file_index = system_index["current_departure_file_index"]
    departure_index_in_file = system_index["departure_index_in_file"]
    ARRIVALS_TIME_INDEX = system_index["arrivals_time_index"]
    DEPARTURES_TIME_INDEX = system_index["departures_time_index"]
    j = system_index["j"]
    arrivals_on_service = system_index["arrivals_on_service"]

    # if last event arrivals, move in arrivals
    if system_index["m_j"] == 0:
        # next arrival inside current file
        if arrivals_index_in_file < system_index["arrivals_covariates_num"] - 1:
            arrivals_index_in_file += 1
            system_index["arrivals_index_in_file"] = arrivals_index_in_file
            arrival_time = system_index["arrivals_covariates_full"][arrivals_index_in_file][ARRIVALS_TIME_INDEX]
            system_index["arrivals_covariate"] = system_index["arrivals_covariates_full"][arrivals_index_in_file]
        else:  # next arrival in new file
            if current_arrival_file_index < system_index["total_num_arrivals_file"]:
                current_arrival_file_index += 1
                system_index["current_arrival_file_index"] = current_arrival_file_index

                arrivals_index_in_file = 0  # start of new file
                system_index["arrivals_index_in_file"] = arrivals_index_in_file
                arrivals_covariates_full = np.loadtxt(
                        data_dir + 'arrivals_covariates_{0}_.txt'.format(current_arrival_file_index))
                system_index["arrivals_covariates_full"] = arrivals_covariates_full
                system_index["arrivals_covariates_num"] = arrivals_covariates_full.shape[0]
                arrival_time = system_index["arrivals_covariates_full"][arrivals_index_in_file][ARRIVALS_TIME_INDEX]
                system_index["arrivals_covariate"] = system_index["arrivals_covariates_full"][arrivals_index_in_file]
            else:
                arrival_time = np.infty
        departure_time = system_index["departure_time"]
        system_index["arrivals_time"] = arrival_time
    # if last event departure move in departures
    if system_index["m_j"] == 1:
        # next departure inside current file
        if departure_index_in_file < system_index["departures_covariates_num"] - 1:
            departure_index_in_file += 1
            system_index["departure_index_in_file"] = departure_index_in_file
            departure_time = system_index["departures_covariates_full"][departure_index_in_file][DEPARTURES_TIME_INDEX]
            system_index["departures_covariate"] = system_index["departures_covariates_full"][arrivals_index_in_file]
        else:  # next arrival in new file
            if current_departure_file_index < system_index["total_num_departures_file"]:
                current_departure_file_index += 1
                system_index["current_departure_file_index"] = current_departure_file_index

                departure_index_in_file = 0  # start of new file
                system_index["departure_index_in_file"] = departure_index_in_file
                departures_covariates_full = np.loadtxt(
                        data_dir + 'departures_covariates_{0}_.txt'.format(current_departure_file_index))
                system_index["departures_covariates_full"] = departures_covariates_full
                system_index["departures_covariates_num"] = departures_covariates_full.shape[0]
                departure_time = system_index["departures_covariates_full"][departure_index_in_file][
                    DEPARTURES_TIME_INDEX]
                system_index["departures_covariate"] = system_index["departures_covariates_full"][
                    departure_index_in_file]
            else:
                departure_time = np.infty
        system_index["departure_time"] = departure_time
        arrival_time = system_index["arrivals_time"]
    # ========================================
    # UPDATE CURRENT TIME
    # ========================================
    if arrival_time < departure_time:
        m_j = 0
        t_j = arrival_time
        arrivals_on_service = arrivals_on_service + 1
    else:
        m_j = 1
        t_j = departure_time
        if arrivals_on_service > 0:
            arrivals_on_service = arrivals_on_service - 1

    j += 1
    system_index["m_j"] = m_j
    system_index["t_j"] = t_j
    system_index["j"] = j
    system_index["arrivals_on_service"] = arrivals_on_service

    return system_index


def next_data_input(input):
    """
    arrival_and_covariates = [arrivals_id (0), accepted_answer_id (1), arrival_time (2), owner_id (3),
                              score (4), favorite_count (5), view_count (6), departure_time(7), tags(5) (7+5)]

    departure_and_covariates = [arrivals_id,arrival_time,parent_question,owner_id,
                                score]

    :param data_dir:
    :param system_index:
    :param U_truncated:
    :return: [1,BACKPROP_SEQ_LENGHT,MAX_SERVICE_LOAD,input_to_training]
    *input_to_training: elapsed_time,
                        covariates,
                        true_remaining_service_time,
                        time_to_next_event (float),
                        leaving_system (bool)
    """

    batch_index = input[0]
    past_t_j = input[1]
    events_batch = input[2]
    batch = input[3]
    system_index = input[4]
    U_truncated = input[5]
    SEQ_LENGHT = input[6]
    ARRIVALS_ID_INDEX = input[7]
    DEPARTURES_ID_INDEX = input[8]
    MAXIMUM_SERVICE_LOAD = input[9]
    ARRIVALS_TIME_INDEX = input[10]
    DEPARTURE_TIME_INDEX = input[11]
    U_DIMENSION = input[12]
    REAL_DEPARTURE_TIME_INDEX = input[13]
    U_SELECTION = input[14]

    system_index['arrivals_time_index'] = ARRIVALS_TIME_INDEX
    system_index['departures_time_index'] = DEPARTURE_TIME_INDEX

    if system_index["t_j"] != np.inf:
        EVENTS_BATCH = []
        BATCH = []
        for j in range(SEQ_LENGHT):
            system_index_next = update_system_index(system_index)
            if np.isnan(system_index["t_j"] - past_t_j):
                left_to_padded = SEQ_LENGHT - len(EVENTS_BATCH)
                U = np.zeros(
                        (left_to_padded, MAXIMUM_SERVICE_LOAD, U_DIMENSION))

                BATCH = np.asarray(BATCH)
                BATCH = np.concatenate([BATCH, U], axis=0)

                EVENTS_BATCH = np.asarray(EVENTS_BATCH)
                non_zero = EVENTS_BATCH.shape[0]
                EVENTS_BATCH = np.concatenate(
                        [EVENTS_BATCH, np.zeros((left_to_padded, 4))], axis=0)
                return batch_index, past_t_j, EVENTS_BATCH, BATCH, system_index, U_truncated, non_zero, \
                       SEQ_LENGHT, ARRIVALS_ID_INDEX, DEPARTURES_ID_INDEX, MAXIMUM_SERVICE_LOAD, ARRIVALS_TIME_INDEX, DEPARTURE_TIME_INDEX, U_DIMENSION, REAL_DEPARTURE_TIME_INDEX, U_SELECTION

            EVENTS_BATCH.append(np.asarray([system_index["t_j"], system_index["t_j"] -
                                            past_t_j, system_index["m_j"], system_index["arrivals_on_service"]]))
            # =========================================================
            # SERVICE DATA UPDATE ACCORDING TO EVENTS
            # =========================================================
            if system_index["m_j"] == 1:  # departure occurred
                if U_truncated is not None:
                    still_on_sytem = np.logical_not(
                            U_truncated[:, ARRIVALS_ID_INDEX] == system_index["departures_covariate"][
                                DEPARTURES_ID_INDEX])
                    # removed arrival from system (departure!)
                    U_truncated = U_truncated[still_on_sytem, :]
            elif system_index["m_j"] == 0:  # arrival ocurred
                if U_truncated is not None:
                    if U_truncated.shape[
                        0] == MAXIMUM_SERVICE_LOAD:  # maximum service load attained, remove oldest arrival
                        U_truncated[:MAXIMUM_SERVICE_LOAD -
                                     1, :] = U_truncated[1:, :]
                        U_truncated[-1, :] = system_index["arrivals_covariate"]
                    else:
                        U_truncated = np.concatenate((U_truncated, np.asarray([system_index["arrivals_covariate"], ])),
                                                     axis=0)
                else:
                    U_truncated = np.asarray(
                            [system_index["arrivals_covariate"], ])
            # =========================================================
            # TIME TO NEXT EVENT
            # =========================================================
            if U_truncated is not None:
                time_to_next_event = np.repeat(system_index_next["t_j"] - system_index["t_j"],
                                               U_truncated.shape[0]).reshape(U_truncated.shape[0], 1)
                leaving_system = (U_truncated[:, ARRIVALS_ID_INDEX] == system_index_next["departures_covariate"][
                    DEPARTURES_ID_INDEX]).reshape(U_truncated.shape[0], 1)
                U = np.concatenate(
                        [U_truncated, time_to_next_event, leaving_system], axis=1)
                U[:, ARRIVALS_TIME_INDEX] = system_index["t_j"] - U_truncated[:,
                                                                  ARRIVALS_TIME_INDEX]  # define elapsed time per arrival

                # define time to real departure (if not departure, time to global censoring time)
                mask = (U[:, REAL_DEPARTURE_TIME_INDEX] < 0)
                sign = np.sign(U[:, REAL_DEPARTURE_TIME_INDEX])
                U[:, REAL_DEPARTURE_TIME_INDEX] = mask * U[:, ARRIVALS_TIME_INDEX] - \
                                                  sign * system_index["t_j"] + \
                                                  U[:, REAL_DEPARTURE_TIME_INDEX]

                left_to_padded = MAXIMUM_SERVICE_LOAD - U_truncated.shape[0]
                # padding guarantess that input data is of same shape
                U = np.pad(U, ((0, left_to_padded), (0, 0)), "constant")
                U = U[:, U_SELECTION]  # remove unwanted covariates ids
                system_index = system_index_next
                past_t_j = system_index["t_j"]
                BATCH.append(U)
            else:
                system_index = system_index_next
                past_t_j = system_index["t_j"]
                U = np.zeros((MAXIMUM_SERVICE_LOAD, U_DIMENSION))
                BATCH.append(U)
        BATCH = np.asarray(BATCH)
        return batch_index, past_t_j, EVENTS_BATCH, BATCH, system_index, U_truncated, len(EVENTS_BATCH), \
               SEQ_LENGHT, ARRIVALS_ID_INDEX, DEPARTURES_ID_INDEX, MAXIMUM_SERVICE_LOAD, ARRIVALS_TIME_INDEX, DEPARTURE_TIME_INDEX, U_DIMENSION, REAL_DEPARTURE_TIME_INDEX, U_SELECTION
    else:
        EVENTS_BATCH = np.zeros((SEQ_LENGHT, 4))
        BATCH = np.zeros((SEQ_LENGHT, MAXIMUM_SERVICE_LOAD, U_DIMENSION))
        return batch_index, system_index["t_j"], EVENTS_BATCH, BATCH, system_index, U_truncated, 0, \
               SEQ_LENGHT, ARRIVALS_ID_INDEX, DEPARTURES_ID_INDEX, MAXIMUM_SERVICE_LOAD, ARRIVALS_TIME_INDEX, DEPARTURE_TIME_INDEX, U_DIMENSION, REAL_DEPARTURE_TIME_INDEX, U_SELECTION


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 = None
        self.ks_d = None
        self._one = tf.constant(1, dtype=tf.compat.v1.int32)
        self.summary = None

        self.best_model_loss = float("inf")
        self.global_step = tf.Variable(
                0, name='global_step_marks', trainable=False)
        self._build_placeholders()
        seed = self._prepare_dirs()

        np.random.seed(seed)
        tf.compat.v1.set_random_seed(seed)

    @abstractmethod
    def _build_placeholders(self):
        return NotImplemented

    @abstractmethod
    def _build_recurrent(self):
        return NotImplemented

    @abstractmethod
    def _build_service_distribution(self):
        return NotImplemented

    @abstractmethod
    def _get_loss(self):
        return NotImplemented

    @abstractmethod
    def _get_discriminator_loss(self):
        return NotImplemented

    @abstractmethod
    def _build_training(self):
        return NotImplemented

    @abstractmethod
    def _get_predicted_marks(self):
        return NotImplemented

    @abstractmethod
    def _build_one_step_rmse(self):
        return NotImplemented

    @abstractmethod
    def train(self, train_set, test_set, sess):
        """
        Args:
            train_data:
            test_data:
            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.compat.v1.sign(tf.compat.v1.reduce_max(tf.compat.v1.abs(sequences), 2))
        sequences_length = tf.compat.v1.reduce_sum(non_zero, 1)
        sequences_length = tf.compat.v1.cast(sequences_length, tf.compat.v1.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 == False:
            os.makedirs(self.checkpoint_dir)
            os.makedirs(self.tensorboard_dir)

        return int(start_time[-3:])

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

    def store_training_epoch(self,
                             session,
                             summary_writer,
                             epoch,
                             iteration,
                             rmse,
                             loss,
                             ks_d,
                             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.assign(rmse),
                self.ks_d.assign(ks_d)
            ])

        else:
            # i.e. when not using GAN models:
            session.run([self.loss.assign(loss),
                         self.rmse.assign(rmse),
                         self.ks_d.assign(ks_d)])

        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")
            logger.info("Saving checkpoint: {} ...".format(model_filename))
            self.saver.save(session, model_filename, epoch)

        self._write_model_parameters(self.checkpoint_dir)

    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, batch_size_test: int = None):
        """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.join(*path.split(os.sep)[:-1])
        if path[0] == os.sep:
            model_directory = os.sep + model_directory
        with open(cls._parameters_file(model_directory)) as f:
            parameters = json.load(f)
        parameters['restore'] = True
        parameters['restore_dir'] = model_directory
        if batch_size_test is not None:
            parameters['batch_size_test'] = batch_size_test
        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 DeepServiceTime(Model):

    def __init__(self, args):
        """
        Args:
            args:
        """
        super(DeepServiceTime, self).__init__(args)
        self._adversarial_models = ("AdversarialServiceTime",
                                    "RecurrentAdversarialServiceTimeBeta",
                                    "RecurrentAdversarialMemPool",
                                    "RecurrentAdversarialTransactions",
                                    "RASBetaStaticCritic",
                                    "SimpleAdversarial")
        self._build_recurrent()
        self._build_service_distribution()
        self._build_training()
        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")

    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('Placeholders'):

            if self.model in ("RecurrentNeuralServiceTime", "RecurrentAdversarialServiceTime"):
                self.input_data = tf.compat.v1.placeholder(
                        tf.compat.v1.float64, [None,
                                     self.args.num_steps,
                                     self.args.maximum_service_load,
                                     self.args.input_dim], 'input_data')
            else:
                self.input_data = tf.compat.v1.placeholder(
                        tf.compat.v1.float64, [None, self.args.num_steps, self.args.input_dim], 'input_data')

            self.input_events = tf.compat.v1.placeholder(tf.compat.v1.float64, [None, self.args.num_steps, self.args.events_dim],
                                               'input_events')

            # seq. length set for dynamic_rnn
            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, self.args.num_steps], 'mask')

            self.T = tf.compat.v1.placeholder(
                    tf.compat.v1.float64, shape=[], name='censoring_time')  # censoring time

            if self.args.use_hidden:

                self.hidden_states_arrivals = tf.compat.v1.placeholder(
                        tf.compat.v1.float64, [None, self.args.num_steps,
                                     self.args.hidden_units],
                        'arrivals_hidden_state')
            else:
                self.hidden_states_arrivals = None

            self.beta_censoring = tf.compat.v1.placeholder(
                    tf.compat.v1.float64, shape=[], name="beta_censoring")

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

            self.training_mode = tf.compat.v1.placeholder(
                    tf.compat.v1.bool, shape=[], name="training_mode")

            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.compat.v1.placeholder(tf.compat.v1.float64, shape=[None, h_dim], name=name)
        h = tf.compat.v1.placeholder(tf.compat.v1.float64, shape=[None, h_dim], name=name)

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

    def _get_loss(self):
        pass

    def _get_discriminator_loss(self):
        pass

    def _build_training(self):
        with tf.compat.v1.name_scope("Train"):

            self.rmse = tf.Variable(
                    float("inf"), False, dtype=tf.compat.v1.float64, name="rmse")

            self.loss = tf.Variable(
                    float("inf"), False, dtype=tf.compat.v1.float64, name="loss")

            self.loss_discriminator = tf.Variable(
                    float("inf"), False, dtype=tf.compat.v1.float64, name="loss_dis")

            self.ks_d = tf.Variable(
                    float("inf"), False, dtype=tf.compat.v1.float64, name="ks_d")

            tf.compat.v1.summary.scalar(self.rmse.op.name, self.rmse)
            tf.compat.v1.summary.scalar(self.loss.op.name, self.loss)
            tf.compat.v1.summary.scalar(self.loss_discriminator.op.name,
                              self.loss_discriminator)
            tf.compat.v1.summary.scalar(self.ks_d.op.name, self.ks_d)

            with tf.compat.v1.name_scope("Losses"):
                if self.model == "NeuralServiceTime":
                    self.mark_loss = self._get_loss()
                else:
                    self.mark_loss = self._get_loss()
                    self.dis_loss = self._get_discriminator_loss()

            self.__build_train_step()

    def __build_train_step(self):

        if self.model in self._adversarial_models:
            optimizer = tf.compat.v1.train.AdamOptimizer(
                    learning_rate=self.args.learning_rate_gan,
                    beta1=0.5,
                    beta2=0.9)

            # Trainable variables within the GENERATOR

            train_var = tf.compat.v1.get_collection(
                    tf.compat.v1.GraphKeys.GLOBAL_VARIABLES, scope="GAN/generator")
            grad = tf.compat.v1.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.compat.v1.get_collection(
                    tf.compat.v1.GraphKeys.GLOBAL_VARIABLES, scope="GAN/discriminator")
            grad_dis = tf.compat.v1.gradients(self.dis_loss, train_dis_var)
            self.train_step_dis = optimizer.apply_gradients(
                    zip(grad_dis, train_dis_var), name="train_step_dis")

        else:

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

            train_var = tf.compat.v1.get_collection(
                    tf.compat.v1.GraphKeys.GLOBAL_VARIABLES, scope="Marks")
            grad = tf.compat.v1.gradients(self.mark_loss, train_var)
            self.train_step = optimizer.apply_gradients(
                    zip(grad, train_var),
                    global_step=self.global_step,
                    name="train_step_marks")

    def _get_predicted_marks(self):
        pass

    def _build_one_step_rmse(self):
        pass

    def train_parallel(self, sess, h_states_arrivals):
        P_DATA_DIR = self.args.parallel_data_dir
        ARRIVALS_TIME_INDEX = self.args.arrivals_time_index
        DEPARTURES_TIME_INDEX = self.args.departures_time_index
        pool = mp.Pool(self.args.number_of_process)
        batch_files = [a for a in os.listdir(P_DATA_DIR) if "tags" not in a]
        NUMBER_OF_BATCHES = len(batch_files)
        batch_directories = [P_DATA_DIR + batch_file +
                             "/" for batch_file in batch_files]
        SEQ_LENGHT = self.args.num_steps
        ARRIVALS_ID_INDEX = self.args.arrivals_id_index
        DEPARTURES_ID_INDEX = self.args.departures_id_index
        MAXIMUM_SERVICE_LOAD = self.args.maximum_service_load
        U_DIMENSION = self.args.u_dimension
        U_SELECTION = self.args.u_selection
        REAL_DEPARTURE_TIME_INDEX = self.args.real_departure_time_index

        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):
                if self.model == "RecurrentNeuralServiceTime":
                    state = sess.run(self.reset_state)
                SYSTEM_INDEX_PER_BATCH = []
                for j, data_dir in enumerate(batch_directories):
                    sys_ix = self.initialize_system_index(
                            data_dir, ARRIVALS_TIME_INDEX, DEPARTURES_TIME_INDEX)
                    SYSTEM_INDEX_PER_BATCH.append(
                            (j, 0., None, None, sys_ix[0], sys_ix[1], None,
                             SEQ_LENGHT,
                             ARRIVALS_ID_INDEX,
                             DEPARTURES_ID_INDEX,
                             MAXIMUM_SERVICE_LOAD,
                             ARRIVALS_TIME_INDEX,
                             DEPARTURES_TIME_INDEX,
                             U_DIMENSION,
                             REAL_DEPARTURE_TIME_INDEX,
                             U_SELECTION))
                sequence_section_index = 0
                seq_ix = 0
                d0 = datetime.datetime.now()

                count_loop = 0
                while True:

                    SYSTEM_INDEX_PER_BATCH = pool.map(
                            next_data_input, SYSTEM_INDEX_PER_BATCH)
                    SYSTEM_INDEX_PER_BATCH.sort()

                    events_data = np.array([batch_line[2]
                                            for batch_line in SYSTEM_INDEX_PER_BATCH])

                    costumers_in_service = np.array([batch_line[3]
                                                     for batch_line in SYSTEM_INDEX_PER_BATCH])

                    padding = np.asarray([batch_line[-1]
                                          for batch_line in SYSTEM_INDEX_PER_BATCH])

                    B, T, L, _ = costumers_in_service.shape

                    s_l = np.ones(B) * T
                    mask = np.ones((B, T, L))
                    h = h_states_arrivals[0, :, seq_ix]

                    if self.model == "RecurrentAdversarialServiceTime":
                        # Repeat data for MC sampling
                        events_data = np.repeat(
                                events_data, self.args.n_mc_samples, axis=0)
                        costumers_in_service = np.repeat(
                                costumers_in_service, 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)
                        h = np.repeat(h, self.args.n_mc_samples, axis=0)

                        if count_loop == 0:
                            # 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)

                    count_loop += 1

                    feed_dict = {self.input_events: events_data,
                                 self.input_data: costumers_in_service,
                                 self.seq_len: s_l,
                                 self.mask: mask,
                                 self.hidden_states_arrivals: h}

                    if self.model == "RecurrentAdversarialServiceTime":
                        feed_dict.update({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.T: 1.0,
                                          self.dropout_prob: self.args.dropout_prob,
                                          self.training_mode: True})

                        for _ in range(self.args.n_discriminator_updates):
                            dis_loss, _, _, state_dis_fake, state_dis_real, state_dis_interpol = sess.run(
                                    [self.dis_loss,
                                     self.train_step_dis,
                                     self.clip_discriminator,
                                     self.next_state_dis_fake,
                                     self.next_state_dis_real,
                                     self.next_state_dis_interpol], feed_dict)

                        for _ in range(self.args.n_generator_updates):
                            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} Step {1}: {2:4f}, {3:4f}".format(epoch,
                                                                                iteration,
                                                                                mark_loss,
                                                                                dis_loss))
                    elif self.model == "RecurrentNeuralServiceTime":
                        mark_loss, _, iteration = sess.run([self.mark_loss, self.train_step, self.global_step],
                                                           feed_dict)

                        if self.logging_interval(iteration, self.args.logging_interval):
                            logger.info("Epoch {0} ({1:.2%}) Step {2}: {3:4f}".format(epoch,
                                                                                      iteration,
                                                                                      iteration,
                                                                                      mark_loss))
                    else:
                        print("ERROR: model not specified")
                        raise Exception

                    sequence_section_index += 1
                    min_t_j = min(
                            [system_index["t_j"] for
                             j, past_t_j, EVENTS_BATCH, BATCH, system_index, U, _, _, _, _, _, _, _
                             in SYSTEM_INDEX_PER_BATCH])
                    if min_t_j == np.infty:
                        df = datetime.datetime.now()
                        print("epoch time: ", (df - d0).total_seconds())
                        break

                self.store_training_epoch(sess, summary_writer,
                                          epoch, iteration, 0, mark_loss)

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

    def initialize_system_index(self, data_dir, ARRIVALS_TIME_INDEX, DEPARTURES_TIME_INDEX):
        """
        :param data_dir:
        :return:
        """
        all_arrivals_files = [a for a in os.listdir(
                data_dir) if "arrivals_covariates" in a]
        arrivals_files_numbers = [int(arrivals_file_name.split(
                "_")[-2]) for arrivals_file_name in all_arrivals_files]
        arrivals_files_numbers.sort()
        all_arrivals_files = ["arrivals_covariates_{0}_.txt".format(
                num) for num in arrivals_files_numbers]
        all_departures_files = [a for a in os.listdir(
                data_dir) if "departures_covariates" in a]
        departures_files_numbers = [int(arrivals_file_name.split(
                "_")[-2]) for arrivals_file_name in all_departures_files]
        departures_files_numbers.sort()
        all_departures_files = ["departures_covariates_{0}_.txt".format(
                num) for num in departures_files_numbers]

        arrivals_covariates_full = np.loadtxt(data_dir + all_arrivals_files[0])
        departures_covariates_full = np.loadtxt(
                data_dir + all_departures_files[0])

        arrival_time = arrivals_covariates_full[0][ARRIVALS_TIME_INDEX]
        departure_time = departures_covariates_full[0][DEPARTURES_TIME_INDEX]
        if arrival_time < departure_time:
            m_j = 0
            t_j = arrival_time
            arrivals_on_service = 1
        else:
            m_j = 1
            t_j = departure_time
            arrivals_on_service = 0

        total_num_arrivals_file = len(all_arrivals_files)
        total_num_departures_file = len(all_departures_files)

        system_index = {"data_dir": data_dir,
                        "current_arrival_file_index": 1,
                        "arrivals_index_in_file": 0,
                        "current_departure_file_index": 1,
                        "departure_index_in_file": 0,
                        "arrivals_time": arrival_time,
                        "arrivals_covariate": arrivals_covariates_full[0],
                        "departure_time": departure_time,
                        "departures_covariate": departures_covariates_full[0],
                        "m_j": m_j,
                        "t_j": t_j,
                        "j": 0,
                        "arrivals_on_service": arrivals_on_service,
                        "arrivals_covariates_full": arrivals_covariates_full,
                        "departures_covariates_full": departures_covariates_full,
                        "arrivals_covariates_num": arrivals_covariates_full.shape[0],
                        "departures_covariates_num": departures_covariates_full.shape[0],
                        "total_num_arrivals_file": total_num_arrivals_file,
                        "total_num_departures_file": total_num_departures_file}

        return system_index, None

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

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

        Returns:

        """
        censoring_time = 1.
        train_data, train_seq_len, train_mask = train_set
        test_data, test_seq_len, test_mask = test_set

        sess.run(self.initialize)  # variable initialization

        summary_writer = self.summary_writer(self.tensorboard_dir, sess)

        NUM_BATCHES, B, NUM_SEQ, _, _ = train_data.shape
        NUM_BATCHES = min(NUM_BATCHES, h_train.shape[0])
        B_T, T_NUM_SEQ, _, _ = test_data.shape
        T_NUM_SEQ = min(T_NUM_SEQ, h_test.shape[1])

        try:
            # TRAINING GAN INDEPENDENTLY FROM RPP:
            TOTAL_ITERATIONS = max(1, self.args.num_epochs * NUM_BATCHES * (
                    NUM_SEQ - 1.))

            for epoch in range(self.args.num_epochs):
                for batch_ix in range(NUM_BATCHES):
                    if self.model == "RecurrentAdversarialServiceTimeBeta":
                        # 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)
                    elif self.model == "RASBetaStaticCritic":
                        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]
                        m = train_mask[batch_ix, :, seq_ix]
                        h = None
                        if h_train is not None:
                            h = h_train[batch_ix, :, seq_ix]
                            h = h[:B, :, :]

                        step = epoch * NUM_BATCHES * NUM_SEQ + batch_ix * NUM_SEQ + seq_ix

                        if self.args.decay_beta_censor == 0:
                            beta_censoring = self.args.censor_penalty
                        else:
                            # exponential building up of censoring term:
                            # 1: (here self.args.decay_beta_censor fixes the percentage strength achieved after all ITE)
                            # decay_rate = - (1.0 / TOTAL_ITERATIONS) * np.log(self.args.decay_beta_censor)
                            # 2: (here self.args.decay_beta_censor is simply a speed factor)
                            decay_rate = - (self.args.decay_beta_censor / TOTAL_ITERATIONS) * np.log(0.01)
                            beta_censoring = self.args.censor_penalty * (1.0 -
                                                                         np.exp(- decay_rate * step))
                            # 3:
                            # percentage = 80.
                            # step_0 = (percentage * TOTAL_ITERATIONS) / 100.
                            # if step < step_0:
                            #     beta_censoring = 0.0
                            # else:
                            #     decay_rate = - (self.args.decay_beta_censor/(TOTAL_ITERATIONS - step_0)) * np.log(0.01)
                            #     beta_censoring = self.args.censor_penalty * (1.0 -
                            #                                                  np.exp(- decay_rate * (step-step_0)))
                        if h_train is not None:
                            feed_dict = {
                                self.input_data: x,
                                self.mask: m,
                                self.hidden_states_arrivals: h,
                                self.beta_censoring: beta_censoring
                            }
                        else:
                            feed_dict = {
                                self.input_data: x,
                                self.mask: m,
                                self.beta_censoring: beta_censoring
                            }

                        if self.model in ("NeuralServiceTime", "MemPool"):
                            mark_loss, _, iteration = sess.run([self.mark_loss, self.train_step, self.global_step],
                                                               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,
                                                                                          mark_loss))

                        elif self.model == "AdversarialServiceTime" or self.model == "SimpleAdversarial":

                            feed_dict.update({self.T: censoring_time,
                                              self.dropout_prob: self.args.dropout_prob,
                                              self.training_mode: True})

                            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, _, iteration = sess.run([self.mark_loss, self.train_step, self.global_step],
                                                                   feed_dict)

                            g1, g2, g3 = sess.run([self.gen_loss_1,
                                                   self.gen_loss_2,
                                                   self.gen_loss_3], feed_dict)

                            if self.logging_interval(iteration, self.args.logging_interval):
                                logger.info(
                                        "Epoch {0} ({1:.2%}) Stp {2}: {3:4f}, {4:4f} LossGen: {5:4f}, {6:4f}, {7:4f} BetaCensor: {8:4f}"
                                            .format(epoch,
                                                    iteration / TOTAL_ITERATIONS,
                                                    iteration, mark_loss, dis_loss, g1, g2, g3, beta_censoring))

                        elif self.model == "RASBetaStaticCritic":

                            feed_dict.update({self.seq_len: s_l,
                                              self.state: state,
                                              self.T: censoring_time,
                                              self.dropout_prob: self.args.dropout_prob,
                                              self.training_mode: True})

                            for _ in range(self.args.n_discriminator_updates):
                                dis_loss, _ = sess.run(
                                        [self.dis_loss, self.train_step_dis], 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} Beta Censoring: {5:4f}"
                                            .format(epoch,
                                                    iteration / TOTAL_ITERATIONS,
                                                    iteration, mark_loss, dis_loss, beta_censoring))

                        elif self.model == "RecurrentAdversarialServiceTimeBeta":

                            feed_dict.update({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.T: censoring_time,
                                              self.dropout_prob: self.args.dropout_prob,
                                              self.training_mode: True})

                            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))

                        else:
                            print("Error: model not defined")
                            raise Exception

                    if self.model == "RecurrentAdversarialServiceTimeBeta" and 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]
                                h = None
                                if h_train is not None:
                                    h = h_train[batch_ix, :, seq_ix]
                                feed_dict.update({self.input_data: x,
                                                  self.mask: m,
                                                  self.hidden_states_arrivals: h,
                                                  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.T: censoring_time,
                                                  self.dropout_prob: self.args.dropout_prob,
                                                  self.training_mode: True})

                                _, 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:
                se = []
                predictions = []
                if self.model == "RecurrentAdversarialServiceTimeBeta" or self.model == "RASBetaStaticCritic":
                    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]
                    h_ = None
                    if h_test is not None:
                        h_ = h_test[:, t_seq_ix]
                    mask = test_mask[:, t_seq_ix]
                    s_l_ = test_seq_len[:, t_seq_ix]

                    if self.model in self._adversarial_models:
                        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)

                    if h_test is not None:
                        feed_dict = {
                            self.input_data: x_,
                            self.hidden_states_arrivals: h_,
                            self.T: censoring_time,
                            self.dropout_prob: 1.0,
                            self.training_mode: False,
                            self.seq_len: s_l_,
                            self.mask: mask
                        }
                    else:
                        feed_dict = {
                            self.input_data: x_,
                            self.T: censoring_time,
                            self.dropout_prob: 1.0,
                            self.training_mode: False,
                            self.seq_len: s_l_,
                            self.mask: mask
                        }

                    if self.model == "RecurrentAdversarialServiceTimeBeta" or self.model == "RASBetaStaticCritic":
                        feed_dict.update({self.state: state})
                        se_seq, prediction, state = sess.run([self.one_step_rmse,
                                                              self.sample_dist,
                                                              self.next_state],
                                                             feed_dict)
                    else:
                        se_seq, prediction = sess.run([self.one_step_rmse,
                                                       self.sample_dist],
                                                      feed_dict)

                    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)
                service_times = np.asarray(predictions)
                ks_d, p = ks_2samp(
                        service_times[:, :, :self.args.n_mc_samples].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)

                if self.model in ("NeuralServiceTime", "MemPool"):
                    self.store_training_epoch(sess, summary_writer,
                                              epoch, iteration, rmse, mark_loss, ks_d)
                else:
                    self.store_training_epoch(sess, summary_writer,
                                              epoch, iteration, rmse, mark_loss, ks_d,
                                              dis_loss)

        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

    def sample(self, test_set, h_test, sess):
        censoring_time = 1.

        test_data, seq_len, test_mask = test_set
        B, NUM_SEQ, _, _ = test_data.shape
        NUM_SEQ = min(NUM_SEQ, h_test.shape[1])

        service_times = []

        if self.model == "RecurrentAdversarialServiceTimeBeta" or self.model == "RASBetaStaticCritic":
            state = self._get_zero_state(B * self.args.n_mc_samples,
                                         self.args.gen_dim)

        for seq_ix in range(NUM_SEQ):
            x_ = test_data[:, seq_ix]
            h_ = None
            if h_test is not None:
                h_ = h_test[:, seq_ix]
                h_ = h_[:B, :, :]
            mask = test_mask[:, seq_ix]
            s_l_ = seq_len[:, seq_ix]

            if self.model in self._adversarial_models:
                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)

            feed_dict = {
                self.input_data: x_,
                self.hidden_states_arrivals: h_,
                self.T: censoring_time,
                self.dropout_prob: 1.0,
                self.training_mode: False,
                self.seq_len: s_l_,
                self.mask: mask
            }

            if self.model == "RecurrentAdversarialServiceTimeBeta" or self.model == "RASBetaStaticCritic":
                feed_dict.update({self.state: state})
                se_seq, prediction, state = sess.run([self.one_step_rmse,
                                                      self.sample_dist,
                                                      self.next_state],
                                                     feed_dict)
            else:
                se_seq, prediction = sess.run([self.one_step_rmse,
                                               self.sample_dist],
                                              feed_dict)

            prediction = np.concatenate(
                    (prediction, test_data[:, seq_ix]), axis=2)  # [B, T, MC + 1]
            service_times.append(prediction)

        service_times = np.concatenate(service_times, axis=1)

        return service_times

class AdversarialServiceTime(DeepServiceTime):

    def __init__(self, args):
        """
        Args:
            args:
        """
        self.model = "AdversarialServiceTime"
        print("Service time model: Adversarial Service Time")
        super(AdversarialServiceTime, self).__init__(args)
        self._build_service_distribution()
        self.sample_dist = self._get_predicted_marks()
        self.saver = tf.compat.v1.train.Saver(max_to_keep=100)

    def _build_recurrent(self):
        pass

    def _build_service_distribution(self):

        h = self.hidden_states_arrivals

        covariates = self.input_data[:, :, :-1]
        # observed service time [B*T, 1]
        s_real = tf.compat.v1.reshape(self.input_data[:, :, -1], [-1, 1])
        s_real_sq = tf.compat.v1.squeeze(s_real)
        with tf.compat.v1.name_scope("build_generator-discriminator"):
            # Indices censored-uncensored data:
            self._index_uncensored = tf.compat.v1.where(
                    tf.compat.v1.greater(s_real_sq, 0), tf.compat.v1.ones_like(s_real_sq), tf.compat.v1.zeros_like(s_real_sq))
            self._index_uncensored = tf.compat.v1.cast(self._index_uncensored, tf.compat.v1.int32)
            self._index_censored = tf.compat.v1.where(tf.compat.v1.less(s_real_sq, 0), tf.compat.v1.zeros_like(s_real_sq), tf.compat.v1.ones_like(s_real_sq))
            self._index_censored = tf.compat.v1.cast(self._index_censored, tf.compat.v1.int32)

            # Generator output
            self._s_gen = self._generator(h, covariates)  # [B*T, 1]

            # Discriminator output (using instance noise: https://arxiv.org/pdf/1610.04490.pdf)

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

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

            # Gradient penalty term:
            self._grad_penalty = self._gradient_penalty_term(s_real, self._s_gen, h, covariates)

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

        with tf.compat.v1.variable_scope("GAN/generator", reuse=tf.compat.v1.AUTO_REUSE):
            with tf.compat.v1.variable_scope("hidden_layers", reuse=tf.compat.v1.AUTO_REUSE):
                shape = tf.compat.v1.shape(h)
                B = shape[0]
                T = shape[1]

                cell_initializer = tf.compat.v1.initializers.glorot_normal()

                # embedding
                # sampling prior normal distr.
                z = tf.compat.v1.random_normal([B * T, self.args.gen_dim], 0.0, 1.0, dtype=tf.compat.v1.float64)
                z = tf.compat.v1.layers.dense(
                        z,
                        self.args.gen_dim,
                        use_bias=False,
                        kernel_initializer=cell_initializer,
                        name="layer_noise_0")

                if self.args.arrivals_hidden_state_bool:
                    h = tf.compat.v1.reshape(
                            h, [-1, self.args.hidden_units])

                    h = tf.compat.v1.layers.dense(
                            h,
                            self.args.gen_dim,
                            kernel_initializer=cell_initializer,
                            name="layer_hidden_0")
                    z = z + h

                if self.args.covariates_bool:
                    covariates = tf.compat.v1.reshape(covariates, [-1, self.args.input_dim - 1])
                    covariates = tf.compat.v1.layers.dense(
                            covariates,
                            self.args.gen_dim,
                            kernel_initializer=cell_initializer,
                            name="layer_emb_cov")
                    z = z + covariates

                # hidden state
                h = tf.compat.v1.nn.leaky_relu(z)  # [B*T, D_GAN]

                if self.args.n_layers_gen > 1:
                    for i in range(self.args.n_layers_gen - 1):
                        h = tf.compat.v1.layers.dense(
                                h,
                                self.args.gen_dim,
                                kernel_initializer=cell_initializer,
                                name="layer_hidden_{}".format(i + 1))
                        if self.args.noisy_layers:
                            z = tf.compat.v1.random_normal(
                                    tf.compat.v1.shape(h), 0.0, 1.0, dtype=tf.compat.v1.float64)
                            z = tf.compat.v1.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.compat.v1.nn.dropout(h, self.dropout_prob)
                        if self.args.batch_normalization:
                            h = tf.compat.v1.layers.batch_normalization(
                                    h, training=self.training_mode)
                        h = tf.compat.v1.nn.leaky_relu(h)  # [B*T, D_GAN]

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

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

    def _discriminator(self, h, covariates, s):
        """
        GAN discriminator
        h (arrival hidden state) [B, T, H]
        covariates [B, T, D-3]
        s (service time) [B*T, 1]
        """
        with tf.compat.v1.variable_scope("GAN/discriminator", reuse=tf.compat.v1.AUTO_REUSE):
            with tf.compat.v1.variable_scope("hidden_layers", reuse=tf.compat.v1.AUTO_REUSE):
                cell_initializer = tf.compat.v1.initializers.glorot_normal()

                # embedding
                s = tf.compat.v1.layers.dense(
                        s,
                        self.args.dis_dim,
                        kernel_initializer=cell_initializer,
                        name="layer_service_time")

                if self.args.arrivals_hidden_state_bool:
                    h = tf.compat.v1.reshape(h, [-1, self.args.hidden_units])
                    h = tf.compat.v1.layers.dense(
                            h,
                            self.args.dis_dim,
                            use_bias=False,
                            kernel_initializer=cell_initializer,
                            name="layer_hidden_0")
                    s = s + h

                if self.args.covariates_bool:
                    covariates = tf.compat.v1.reshape(covariates, [-1, self.args.input_dim - 1])
                    covariates = tf.compat.v1.layers.dense(
                            covariates,
                            self.args.dis_dim,
                            use_bias=False,
                            kernel_initializer=cell_initializer,
                            name="layer_emb_cov")
                    s = s + covariates

                # hidden state
                h = tf.compat.v1.nn.leaky_relu(s)  # [B*T, D_GAN]

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

            with tf.compat.v1.variable_scope(
                    "discriminator_output", reuse=tf.compat.v1.AUTO_REUSE):
                h = tf.compat.v1.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.compat.v1.name_scope("gradient_penalty_term"):
            eta = tf.compat.v1.random_uniform(
                    tf.compat.v1.shape(s_real), minval=0.0, maxval=0.1, dtype=tf.compat.v1.float64)

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

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

            return tf.compat.v1.abs(grad_penalty)

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

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

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

                dis_loss = tf.compat.v1.reshape(dis_loss, [batch_size, seq_length])
                dis_loss = tf.compat.v1.multiply(dis_loss, self.mask)
                dis_loss = tf.compat.v1.reshape(dis_loss, [-1, 1])

                # dis_loss = tf.compat.v1.squeeze(
                #     tf.compat.v1.gather(dis_loss, self._index_uncensored))

                dis_loss = tf.compat.v1.dynamic_partition(tf.compat.v1.squeeze(dis_loss), self._index_uncensored, 2)[1]

                dis_loss = tf.compat.v1.reduce_sum(dis_loss)

            return dis_loss

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

                gen_loss_1 = -self._logit_fake

                gen_loss_1 = tf.compat.v1.reshape(gen_loss_1, [batch_size, seq_length])
                gen_loss_1 = tf.compat.v1.multiply(gen_loss_1, self.mask)
                gen_loss_1 = tf.compat.v1.reshape(gen_loss_1, [-1, 1])

                # gen_loss_1 = tf.compat.v1.squeeze(
                #     tf.compat.v1.gather(gen_loss_1, self._index_uncensored))

                _, gen_loss_1 = tf.compat.v1.dynamic_partition(tf.compat.v1.squeeze(gen_loss_1),
                                                     self._index_uncensored,
                                                     2)

                h = self.hidden_states_arrivals
                covariates = self.input_data[:, :, :-1]
                s_gen = self._generator(h, covariates)

                # regulariser for censored events

                gen_loss_2 = tf.compat.v1.maximum(
                        tf.compat.v1.cast(0.0, tf.compat.v1.float64), self.T - s_gen)
                gen_loss_2 = tf.compat.v1.reshape(gen_loss_2, [batch_size, seq_length])
                gen_loss_2 = tf.compat.v1.multiply(gen_loss_2, self.mask)
                gen_loss_2 = tf.compat.v1.reshape(gen_loss_2, [-1, 1])

                # gen_loss_2 = tf.compat.v1.squeeze(
                #     tf.compat.v1.gather(gen_loss_2, self._index_censored))

                _, gen_loss_2 = tf.compat.v1.dynamic_partition(tf.compat.v1.squeeze(gen_loss_2),
                                                     self._index_censored,
                                                     2)

                # regulariser for uncensored events

                s_real = tf.compat.v1.reshape(self.input_data[:, :, -1], [-1, 1])
                gen_loss_3 = tf.compat.v1.abs(s_real - s_gen)
                gen_loss_3 = tf.compat.v1.reshape(gen_loss_3, [batch_size, seq_length])
                gen_loss_3 = tf.compat.v1.multiply(gen_loss_3, self.mask)
                gen_loss_3 = tf.compat.v1.reshape(gen_loss_3, [-1, 1])

                # gen_loss_3 = tf.compat.v1.squeeze(
                #     tf.compat.v1.gather(gen_loss_3, self._index_uncensored))

                _, gen_loss_3 = tf.compat.v1.dynamic_partition(tf.compat.v1.squeeze(gen_loss_3),
                                                     self._index_uncensored,
                                                     2)

                gen_loss = tf.compat.v1.reduce_sum(
                        gen_loss_1) + self.beta_censoring * tf.compat.v1.reduce_sum(gen_loss_2) + tf.compat.v1.reduce_sum(gen_loss_3)

                self.gen_loss_1 = tf.compat.v1.reduce_sum(gen_loss_1)
                self.gen_loss_2 = tf.compat.v1.reduce_sum(gen_loss_2)
                self.gen_loss_3 = tf.compat.v1.reduce_sum(gen_loss_3)

            return gen_loss

    def _predict_mark(self, output, covariates):
        """
        Args:
            output:
            covariates:
        """
        shape = tf.compat.v1.shape(output)
        next_mark = self._generator(output, covariates)
        next_mark = tf.compat.v1.reshape(next_mark, (shape[0], shape[1]))

        return next_mark

    def _get_predicted_marks(self):

        covariates = self.input_data[:, :, :-1]
        shape = tf.compat.v1.shape(self.hidden_states_arrivals)

        next_mark = self._generator(self.hidden_states_arrivals,
                                    covariates)  # [B*MC*T, 1]

        next_mark = tf.compat.v1.reshape(next_mark,
                               (shape[0], shape[1]))  # [B*MC, T]

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

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

        return next_mark

    def _build_one_step_rmse(self):

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

        marks_predicted = self._predict_mark(
                self.hidden_states_arrivals, covariates)

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

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

        # Uncensored error:
        uncensored_error = tf.compat.v1.abs(marks_predicted - marks_real)

        # Censored error:
        censored_error = tf.compat.v1.maximum(tf.compat.v1.zeros_like(
                marks_predicted), self.T - marks_predicted)

        error = tf.compat.v1.where(tf.compat.v1.less(marks_real, 0),
                         censored_error, uncensored_error)

        self.one_step_rmse = error


class RecurrentAdversarialServiceTimeBeta(DeepServiceTime):

    def __init__(self, args):
        """
        Args:
            args:
        """
        self.model = "RecurrentAdversarialServiceTimeBeta"
        super(RecurrentAdversarialServiceTimeBeta, self).__init__(args)
        self._build_service_distribution()
        self.sample_dist = self._get_predicted_marks()
        self.saver = tf.compat.v1.train.Saver(max_to_keep=100)

    def _build_recurrent(self):
        pass

    def _build_service_distribution(self):
        self._build_recurrent_adversarial()

    def _build_recurrent_adversarial(self):
        h = self.hidden_states_arrivals
        covariates = self.input_data[:, :, :-1]

        s_real = tf.compat.v1.reshape(self.input_data[:, :, -1],
                            [-1, 1])  # observed service time [B*T, 1]
        s_real_sq = tf.compat.v1.squeeze(s_real)
        with tf.compat.v1.name_scope("build_generator-discriminator"):
            # Indices censored-uncensored data:
            self._index_uncensored = tf.compat.v1.where(
                    tf.compat.v1.greater(s_real_sq, 0), tf.compat.v1.ones_like(s_real_sq), tf.compat.v1.zeros_like(s_real_sq))
            self._index_uncensored = tf.compat.v1.cast(self._index_uncensored, tf.compat.v1.int32)
            self._index_censored = tf.compat.v1.where(tf.compat.v1.less(s_real_sq, 0), tf.compat.v1.zeros_like(s_real_sq), tf.compat.v1.ones_like(s_real_sq))
            self._index_censored = tf.compat.v1.cast(self._index_censored, tf.compat.v1.int32)
            # Generator output
            self._s_gen, self.next_state = self._generator(
                    h, covariates, self.state)  # [B*T, 1], [B, T, H]

            # Discriminator output
            self._logit_real, self.next_state_dis_real = self._discriminator(
                    h, covariates, s_real, self.state_dis_real)  # [B*T, 1]

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

            self._grad_penalty, self.next_state_dis_interpol = self._gradient_penalty_term(
                    s_real, self._s_gen, h, covariates)

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

        with tf.compat.v1.variable_scope("GAN/generator", reuse=tf.compat.v1.AUTO_REUSE):
            with tf.compat.v1.variable_scope("hidden_layers", reuse=tf.compat.v1.AUTO_REUSE):
                h = tf.compat.v1.reshape(
                        h, [-1, self.args.hidden_units])

                # embedding
                cell_initializer = tf.compat.v1.initializers.glorot_normal()
                h = tf.compat.v1.layers.dense(
                        h,
                        self.args.gen_dim,
                        kernel_initializer=cell_initializer,
                        name="layer_hidden_0")

                if self.args.covariates_bool:
                    covariates = tf.compat.v1.reshape(covariates, [-1, self.args.input_dim - 1])
                    covariates = tf.compat.v1.layers.dense(
                            covariates,
                            self.args.gen_dim,
                            kernel_initializer=cell_initializer,
                            name="layer_emb_cov")
                    h = h + covariates

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

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

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

                h = tf.compat.v1.stack(h)
                h = tf.compat.v1.transpose(h, [1, 0, 2])
                h = tf.compat.v1.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.compat.v1.layers.dense(
                                h,
                                self.args.gen_dim,
                                kernel_initializer=cell_initializer,
                                name="layer_hidden_{}".format(i + 1))
                        if self.args.noisy_layers:
                            z = tf.compat.v1.random_normal(
                                    tf.compat.v1.shape(h), 0.0, 1.0, dtype=tf.compat.v1.float64)
                            z = tf.compat.v1.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.compat.v1.nn.dropout(h, self.dropout_prob)
                        if self.args.batch_normalization:
                            h = tf.compat.v1.layers.batch_normalization(
                                    h, training=self.training_mode)
                        h = tf.compat.v1.nn.leaky_relu(h)  # [B*T, D_GAN]

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

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

    def _discriminator(self, h, covariates, s, state):
        """
        GAN discriminator
        h (arrival hidden state) [B, T, H]
        covariates [B, T, D-3]
        s (service time) [B*T, 1]
        """
        with tf.compat.v1.variable_scope("GAN/discriminator", reuse=tf.compat.v1.AUTO_REUSE):
            with tf.compat.v1.variable_scope("hidden_layers", reuse=tf.compat.v1.AUTO_REUSE):
                # reshaping, etc
                h = tf.compat.v1.reshape(
                        h, [-1, self.args.hidden_units])
                s = tf.compat.v1.reshape(s, [-1, 1])

                # embedding
                # cell_initializer = tf.keras.initializers.he_normal(seed=None)
                cell_initializer = tf.compat.v1.initializers.glorot_normal()

                h = tf.compat.v1.layers.dense(
                        h,
                        self.args.dis_dim,
                        kernel_initializer=cell_initializer,
                        name="layer_hidden_0")

                if self.args.covariates_bool:
                    covariates = tf.compat.v1.reshape(covariates, [-1, self.args.input_dim - 1])
                    covariates = tf.compat.v1.layers.dense(
                            covariates,
                            self.args.dis_dim,
                            kernel_initializer=cell_initializer,
                            name="layer_emb_cov")
                    h = h + covariates

                s = tf.compat.v1.layers.dense(
                        s,
                        self.args.dis_dim,
                        use_bias=False,
                        kernel_initializer=cell_initializer,
                        name="layer_service_time")

                # hidden state
                h = tf.compat.v1.nn.leaky_relu(h + s)  # [B*T, D_GAN]
                h = tf.compat.v1.reshape(
                        h, [-1, self.args.num_steps, self.args.dis_dim])
                h = tf.compat.v1.transpose(h, [1, 0, 2])
                h = tf.compat.v1.unstack(h)

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

                h = tf.compat.v1.stack(h)
                h = tf.compat.v1.transpose(h, [1, 0, 2])
                h = tf.compat.v1.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.compat.v1.layers.dense(
                                h,
                                self.args.dis_dim,
                                kernel_initializer=cell_initializer,
                                name="layer_hidden_{}".format(i + 1))
                        if self.args.dropout:
                            h = tf.compat.v1.nn.dropout(h, self.dropout_prob)
                        if self.args.batch_normalization:
                            h = tf.compat.v1.layers.batch_normalization(
                                    h, training=self.training_mode)
                        h = tf.compat.v1.nn.leaky_relu(h)  # [B*T, D_GAN]

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

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

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

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

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

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

            return tf.compat.v1.abs(grad_penalty), next_state_interpol

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

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

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

                dis_loss = tf.compat.v1.reshape(dis_loss, [batch_size, seq_length])
                dis_loss = tf.compat.v1.multiply(dis_loss, self.mask)
                dis_loss = tf.compat.v1.reshape(dis_loss, [-1, 1])

                # dis_loss = tf.compat.v1.squeeze(
                #     tf.compat.v1.gather(dis_loss, self._index_uncensored))

                _, dis_loss = tf.compat.v1.dynamic_partition(tf.compat.v1.squeeze(dis_loss),
                                                   self._index_uncensored,
                                                   2)

                dis_loss = tf.compat.v1.reduce_sum(dis_loss)

            return dis_loss

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

                gen_loss_1 = -self._logit_fake

                gen_loss_1 = tf.compat.v1.reshape(gen_loss_1, [batch_size, seq_length])
                gen_loss_1 = tf.compat.v1.multiply(gen_loss_1, self.mask)
                gen_loss_1 = tf.compat.v1.reshape(gen_loss_1, [-1, 1])

                # gen_loss_1 = tf.compat.v1.squeeze(
                #     tf.compat.v1.gather(gen_loss_1, self._index_uncensored))

                _, gen_loss_1 = tf.compat.v1.dynamic_partition(tf.compat.v1.squeeze(gen_loss_1),
                                                     self._index_uncensored,
                                                     2)

                s_gen = self._s_gen

                # regulariser for censored events
                gen_loss_2 = tf.compat.v1.maximum(
                        tf.compat.v1.cast(0.0, tf.compat.v1.float64), self.T - s_gen)
                gen_loss_2 = tf.compat.v1.reshape(gen_loss_2, [batch_size, seq_length])
                gen_loss_2 = tf.compat.v1.multiply(gen_loss_2, self.mask)
                gen_loss_2 = tf.compat.v1.reshape(gen_loss_2, [-1, 1])

                # gen_loss_2 = tf.compat.v1.squeeze(
                #     tf.compat.v1.gather(gen_loss_2, self._index_censored))

                _, gen_loss_2 = tf.compat.v1.dynamic_partition(tf.compat.v1.squeeze(gen_loss_2),
                                                     self._index_censored,
                                                     2)

                # regulariser for uncensored events

                s_real = tf.compat.v1.reshape(self.input_data[:, :, -1], [-1, 1])
                gen_loss_3 = tf.compat.v1.abs(s_real - s_gen)
                gen_loss_3 = tf.compat.v1.reshape(gen_loss_3, [batch_size, seq_length])
                gen_loss_3 = tf.compat.v1.multiply(gen_loss_3, self.mask)
                gen_loss_3 = tf.compat.v1.reshape(gen_loss_3, [-1, 1])

                # gen_loss_3 = tf.compat.v1.squeeze(
                #     tf.compat.v1.gather(gen_loss_3, self._index_uncensored))

                _, gen_loss_3 = tf.compat.v1.dynamic_partition(tf.compat.v1.squeeze(gen_loss_3),
                                                     self._index_uncensored,
                                                     2)

                gen_loss = tf.compat.v1.reduce_sum(gen_loss_1) + self.beta_censoring * tf.compat.v1.reduce_sum(
                        gen_loss_2) + tf.compat.v1.reduce_sum(gen_loss_3)

                self.gen_loss_1 = tf.compat.v1.reduce_sum(gen_loss_1)
                self.gen_loss_2 = tf.compat.v1.reduce_sum(gen_loss_2)
                self.gen_loss_3 = tf.compat.v1.reduce_sum(gen_loss_3)

            return gen_loss

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

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

        return next_mark, next_state

    def _get_predicted_marks(self):

        covariates = self.input_data[:, :, :-1]
        shape = tf.compat.v1.shape(self.hidden_states_arrivals)

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

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

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

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

        return next_mark

    def _build_one_step_rmse(self):

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

        marks_predicted, _ = self._predict_mark(
                self.hidden_states_arrivals, covariates)

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

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

        # Uncensored error:
        uncensored_error = tf.compat.v1.abs(marks_predicted - marks_real)

        # Censored error:
        censored_error = tf.compat.v1.maximum(tf.compat.v1.zeros_like(
                marks_predicted), self.T - marks_predicted)

        error = tf.compat.v1.where(tf.compat.v1.less(marks_real, 0),
                         censored_error, uncensored_error)

        self.one_step_rmse = error

class RecurrentAdversarialServiceTime(DeepServiceTime):

    def __init__(self, args):
        """
        Args:
            args:
        """
        self.model = "RecurrentAdversarialServiceTime"
        super(RecurrentAdversarialServiceTime, self).__init__(args)
        self._build_service_distribution()
        self.saver = tf.compat.v1.train.Saver(max_to_keep=100)

    def _build_service_distribution(self):
        if self.args.GAN_model != "WGAN":
            print(
                "Improved WGAN is not possible within TF: 2nd order grads. problem"
            )
            print("Standard GAN is not yet implemented")
            raise Exception
        else:
            self._build_recurrent_adversarial()

    def _build_recurrent_adversarial(self):
        event_data = self.input_data
        service_data = self.input_data[:, :, :, :-3]

        s_remaining_real = self.input_data[:, :, :, -3]
        s_remaining_real = tf.compat.v1.reshape(s_remaining_real, [-1, 1])

        # [B, T, U, 1]  if one(1) it's (locally) uncensored
        uncensoring_bool = self.input_data[:, :, :, -1]
        uncensoring_bool = tf.compat.v1.reshape(uncensoring_bool, [-1, 1])  # [B*T*U, 1]

        self._index_uncensored = tf.compat.v1.where(tf.compat.v1.squeeze(uncensoring_bool))

        self._index_censored_local = tf.compat.v1.where(
                tf.compat.v1.squeeze((-1) * uncensoring_bool + 1))
        self._index_censored_global = tf.compat.v1.where(
                tf.compat.v1.less(tf.compat.v1.squeeze(s_remaining_real), 0.))

        with tf.compat.v1.name_scope("build_generator-discriminator"):
            # Generator output
            self._s_gen, self.next_state = self._generator(service_data,
                                                           event_data,
                                                           self.state)  # [B*T*U, 1],[B, T, H]

            # Discriminator output
            self._logit_real, self.next_state_dis_real = self._discriminator(service_data,
                                                                             event_data,
                                                                             s_remaining_real,
                                                                             self.state_dis_real)

            self._logit_fake, self.next_state_dis_fake = self._discriminator(service_data,
                                                                             event_data,
                                                                             self._s_gen,
                                                                             self.state_dis_fake)

    def _generator(self, service_data, event_data, state):
        """
        GAN generator
        service_data: concat(elapsed_time, cov) [B, T, U, D]
        event_data: concat(arrival_time, interarrival_time, event_type) [B, T, D']
        state: rnn state [B, D_GAN]
        """

        with tf.compat.v1.variable_scope("GAN/generator", reuse=tf.compat.v1.AUTO_REUSE):
            with tf.compat.v1.variable_scope("hidden_layers", reuse=tf.compat.v1.AUTO_REUSE):
                shape = tf.compat.v1.shape(service_data)
                batch_size = shape[0]
                seq_lenght = shape[1]

                cell_initializer = tf.compat.v1.initializers.glorot_normal()

                # embedding_1: event_data
                h = tf.compat.v1.reshape(
                        event_data, [-1, self.args.events_dim])  # [B*T, D']
                h = tf.compat.v1.layers.dense(h, self.args.gen_dim,
                                    kernel_initializer=cell_initializer,
                                    name="layer_hidden_event_data_0")

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

                # hidden state
                h = tf.compat.v1.nn.leaky_relu(h + z)  # [B*T, D_GAN]
                h = tf.compat.v1.reshape(h, [batch_size, seq_lenght, self.args.gen_dim])

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

                h = tf.compat.v1.tile(tf.compat.v1.expand_dims(h, axis=2), [
                    1, 1, self.args.maximum_service_load, 1])  # [B, T, U, D_GAN]

                # [B, T, U, D_GAN + D]  with D = input_dim - 3
                h = tf.compat.v1.concat([h, service_data], axis=-1)
                h = tf.compat.v1.reshape(h,
                               [-1, self.args.gen_dim + self.args.input_dim - 3])

                # embedding_2: service_data + hidden
                h = tf.compat.v1.layers.dense(h,
                                    self.args.gen_dim,
                                    kernel_initializer=cell_initializer,
                                    name="layer_hidden_service_data_0")

                # sampling_2:
                z = tf.compat.v1.random_normal(tf.compat.v1.shape(h), 0.0, 1.0, dtype=tf.compat.v1.float64)
                z = tf.compat.v1.layers.dense(z, self.args.gen_dim,
                                    use_bias=False,
                                    kernel_initializer=cell_initializer,
                                    name="layer_noise_service_data_0")

                h = h + z

                if self.args.dropout:
                    h = tf.compat.v1.nn.dropout(h, self.dropout_prob)
                if self.args.batch_normalization:
                    h = tf.compat.v1.layers.batch_normalization(
                            h, training=self.training_mode)

                h = tf.compat.v1.nn.leaky_relu(h)  # [B*T*U, D_GAN]

                if self.args.n_layers_gen > 1:
                    for i in range(self.args.n_layers_gen - 1):
                        h = tf.compat.v1.layers.dense(
                                h,
                                self.args.gen_dim,
                                kernel_initializer=cell_initializer,
                                name="layer_hidden_{}".format(i + 1))
                        if self.args.noisy_layers:
                            z = tf.compat.v1.random_normal(
                                    tf.compat.v1.shape(h), 0.0, 1.0, dtype=tf.compat.v1.float64)
                            z = tf.compat.v1.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.compat.v1.nn.dropout(h, self.dropout_prob)
                        if self.args.batch_normalization:
                            h = tf.compat.v1.layers.batch_normalization(
                                    h, training=self.training_mode)
                        h = tf.compat.v1.nn.leaky_relu(h)  # [B*T*U, D_GAN]

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

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

    def _discriminator(self, service_data, event_data, s, state):
        """
        GAN discriminator
        service_data: concat(elapsed_time, cov) [B, T, U, D]
        event_data: concat(arrival_time, interarrival_time, event_type) [B, T, D']
        s (service time) [B*T*U, 1]
        state: rnn state [B, D_GAN]
        """
        with tf.compat.v1.variable_scope("GAN/discriminator", reuse=tf.compat.v1.AUTO_REUSE):
            with tf.compat.v1.variable_scope("hidden_layers", reuse=tf.compat.v1.AUTO_REUSE):
                shape = tf.compat.v1.shape(service_data)
                batch_size = shape[0]
                seq_lenght = shape[1]

                cell_initializer = tf.compat.v1.initializers.glorot_normal()

                # embedding_1: event_data
                h = tf.compat.v1.reshape(event_data, [-1, self.args.events_dim])
                h = tf.compat.v1.layers.dense(h, self.args.dis_dim,
                                    kernel_initializer=cell_initializer,
                                    name="layer_hidden_event_data_0")
                h = tf.compat.v1.nn.leaky_relu(h)  # [B*T, D_GAN]
                h = tf.compat.v1.reshape(h, [batch_size, seq_lenght, self.args.dis_dim])

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

                h = tf.compat.v1.tile(tf.compat.v1.expand_dims(h, axis=2), [
                    1, 1, self.args.maximum_service_load, 1])  # [B, T, U, D_GAN]

                # [h, T, U, D_GAN + D]  with D = input_dim - 3
                h = tf.compat.v1.concat([h, service_data], axis=-1)
                h = tf.compat.v1.reshape(h,
                               [-1, self.args.dis_dim + self.args.input_dim - 3])

                # embedding_2: service_data + hidden
                h = tf.compat.v1.layers.dense(h,
                                    self.args.dis_dim,
                                    kernel_initializer=cell_initializer,
                                    name="layer_hidden_service_data_0")

                # embedding_3: remaining time
                s = tf.compat.v1.layers.dense(s, self.args.dis_dim,
                                    use_bias=False,
                                    kernel_initializer=cell_initializer,
                                    name="layer_service_time")

                h = h + s
                if self.args.dropout:
                    h = tf.compat.v1.nn.dropout(h, self.dropout_prob)
                if self.args.batch_normalization:
                    h = tf.compat.v1.layers.batch_normalization(
                            h, training=self.training_mode)
                h = tf.compat.v1.nn.leaky_relu(h)  # [B*T*U, D_GAN]

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

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

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

    def _get_discriminator_loss(self):
        with tf.compat.v1.name_scope("adversarial_loss"):
            with tf.compat.v1.name_scope("discriminator_component"):
                shape = tf.compat.v1.shape(self.input_data)
                batch_size = shape[0]
                seq_length = shape[1]
                maximum_service_load = shape[2]

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

                dis_loss = tf.compat.v1.reshape(
                        dis_loss, [batch_size, seq_length, maximum_service_load])
                dis_loss = tf.compat.v1.multiply(dis_loss, self.mask)
                dis_loss = tf.compat.v1.reshape(dis_loss, [-1, 1])
                dis_loss = tf.compat.v1.squeeze(
                        tf.compat.v1.gather(dis_loss, self._index_uncensored))

                dis_loss = tf.compat.v1.reduce_sum(dis_loss)

            return dis_loss

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

                gen_loss_1 = -self._logit_fake

                gen_loss_1 = tf.compat.v1.reshape(
                        gen_loss_1, [batch_size, seq_length, maximum_service_load])
                gen_loss_1 = tf.compat.v1.multiply(gen_loss_1, self.mask)
                gen_loss_1 = tf.compat.v1.reshape(gen_loss_1, [-1, 1])
                gen_loss_1 = tf.compat.v1.squeeze(
                        tf.compat.v1.gather(gen_loss_1, self._index_uncensored))

                s_gen = self._s_gen

                # regulariser for global censored events

                gen_loss_2 = tf.compat.v1.maximum(
                        tf.compat.v1.cast(0.0, tf.compat.v1.float64), self.T - s_gen)
                gen_loss_2 = tf.compat.v1.reshape(
                        gen_loss_2, [batch_size, seq_length, maximum_service_load])
                gen_loss_2 = tf.compat.v1.multiply(gen_loss_2, self.mask)
                gen_loss_2 = tf.compat.v1.reshape(gen_loss_2, [-1, 1])
                gen_loss_2 = tf.compat.v1.squeeze(
                        tf.compat.v1.gather(gen_loss_2, self._index_censored_global))

                # regulariser for local censored events

                time_to_next_event = self.input_data[:, :, :, -2]
                time_to_next_event = tf.compat.v1.reshape(time_to_next_event, [-1, 1])

                gen_loss_2 = tf.compat.v1.maximum(
                        tf.compat.v1.cast(0.0, tf.compat.v1.float64), time_to_next_event - s_gen)
                gen_loss_2 = tf.compat.v1.reshape(
                        gen_loss_2, [batch_size, seq_length, maximum_service_load])
                gen_loss_2 = tf.compat.v1.multiply(gen_loss_2, self.mask)
                gen_loss_2 = tf.compat.v1.reshape(gen_loss_2, [-1, 1])
                gen_loss_2 = tf.compat.v1.squeeze(
                        tf.compat.v1.gather(gen_loss_2, self._index_censored_local))

                # regulariser for uncensored events

                s_remaining_real = self.input_data[:, :, :, -3]
                s_remaining_real = tf.compat.v1.reshape(s_remaining_real, [-1, 1])

                gen_loss_3 = tf.compat.v1.abs(s_remaining_real - s_gen)
                gen_loss_3 = tf.compat.v1.reshape(
                        gen_loss_3, [batch_size, seq_length, maximum_service_load])
                gen_loss_3 = tf.compat.v1.multiply(gen_loss_3, self.mask)
                gen_loss_3 = tf.compat.v1.reshape(gen_loss_3, [-1, 1])
                gen_loss_3 = tf.compat.v1.squeeze(
                        tf.compat.v1.gather(gen_loss_3, self._index_uncensored))

                gen_loss = tf.compat.v1.reduce_sum(gen_loss_1) + tf.compat.v1.reduce_sum(gen_loss_2) \
                           + tf.compat.v1.reduce_sum(gen_loss_3)

                self.gen_loss_1 = tf.compat.v1.reduce_sum(gen_loss_1)
                self.gen_loss_2 = tf.compat.v1.reduce_sum(gen_loss_2)
                self.gen_loss_3 = tf.compat.v1.reduce_sum(gen_loss_3)

            return gen_loss

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

        # service_data, event_data, state):

        shape = tf.compat.v1.shape(service_data)
        next_mark, next_state = self._generator(
                service_data, event_data, self.state)
        next_mark = tf.compat.v1.reshape(next_mark, (shape[0], shape[1], shape[2]))

        return next_mark, next_state

    def _build_one_step_rmse(self):

        event_data = self.input_data
        service_data = self.input_data[:, :, :, :-3]
        marks_predicted, next_state = self._predict_mark(
                service_data, event_data)

        s_remaining_real = tf.compat.v1.squeeze(
                self.input_data[:, :, :, -3])  # [B, T, U]

        uncensored_error = tf.compat.v1.abs(marks_predicted - s_remaining_real)

        time_to_next_event = tf.compat.v1.squeeze(self.input_data[:, :, :, -2])
        censored_error_local = tf.compat.v1.maximum(tf.compat.v1.zeros_like(marks_predicted),
                                          time_to_next_event - marks_predicted)

        # if one(1) it's (locally) uncensored
        uncensoring_bool = tf.compat.v1.squeeze(self.input_data[:, :, :, -1])

        error_local = tf.compat.v1.where(
                uncensoring_bool, uncensored_error, censored_error_local)

        censored_error_global = tf.compat.v1.maximum(tf.compat.v1.zeros_like(marks_predicted),
                                           self.T - marks_predicted)

        error_global = tf.compat.v1.where(
                tf.compat.v1.less(s_remaining_real, 0), censored_error_global, 0.)

        self.one_step_rmse = tf.compat.v1.reduce_mean(error_local + error_global)

        return next_state


class SimpleAdversarial(DeepServiceTime):

    def __init__(self, args):
        """
        Args:
            args:
        """
        self.model = "SimpleAdversarial"
        print("Service time model: Simple adversarial")
        super(SimpleAdversarial, self).__init__(args)
        self._build_service_distribution()
        self.sample_dist = self._get_predicted_marks()
        self.saver = tf.compat.v1.train.Saver(max_to_keep=100)

    def _build_recurrent(self):
        pass

    def _build_service_distribution(self):
        # observed service time [B*T, 1]
        s_real = tf.compat.v1.reshape(self.input_data[:, :, -1], [-1, 1])

        with tf.compat.v1.name_scope("build_generator-discriminator"):
            # Generator output
            self._s_gen = self._generator()  # [B*T, 1]

            self._logit_real = self._discriminator(s_real)  # [B*T, 1]

            self._logit_fake = self._discriminator(self._s_gen)  # [B*T, 1]

            if self.args.GAN_model == "improved_WGAN":
                self._grad_penalty = self._gradient_penalty_term(
                        s_real, self._s_gen)

    def _generator(self):
        """
        GAN generator
        h (arrival hidden state) [B, T, H]
        covariates [B, T, D-1]
        """
        with tf.compat.v1.variable_scope("GAN/generator", reuse=tf.compat.v1.AUTO_REUSE):
            with tf.compat.v1.variable_scope("hidden_layers", reuse=tf.compat.v1.AUTO_REUSE):
                # sampling prior normal distr.
                cell_initializer = tf.compat.v1.initializers.glorot_normal()
                shape = tf.compat.v1.shape(self.input_data)

                z = tf.compat.v1.random_normal(
                        [shape[0] * shape[1], self.args.gen_dim], 0.0, 1.0, dtype=tf.compat.v1.float64)
                z = tf.compat.v1.layers.dense(
                        z,
                        self.args.gen_dim,
                        use_bias=False,
                        kernel_initializer=cell_initializer,
                        name="layer_noise_0")

                h = tf.compat.v1.nn.leaky_relu(z)  # [B*T, D_GAN]
                if self.args.n_layers_gen > 1:
                    for i in range(self.args.n_layers_gen - 1):
                        h = tf.compat.v1.layers.dense(
                                h,
                                self.args.gen_dim,
                                kernel_initializer=cell_initializer,
                                name="layer_hidden_{}".format(i + 1))
                        if self.args.noisy_layers:
                            z = tf.compat.v1.random_normal(
                                    tf.compat.v1.shape(h), 0.0, 1.0, dtype=tf.compat.v1.float64)
                            z = tf.compat.v1.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.compat.v1.nn.dropout(h, self.dropout_prob)
                        if self.args.batch_normalization:
                            h = tf.compat.v1.layers.batch_normalization(
                                    h, training=self.training_mode)
                        h = tf.compat.v1.nn.leaky_relu(h)  # [B*T, D_GAN]

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

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

    def _discriminator(self, s):
        """
        GAN discriminator
        h (arrival hidden state) [B, T, H]
        covariates [B, T, D-3]
        s (service time) [B*T, 1]
        """
        with tf.compat.v1.variable_scope("GAN/discriminator", reuse=tf.compat.v1.AUTO_REUSE):
            with tf.compat.v1.variable_scope("hidden_layers", reuse=tf.compat.v1.AUTO_REUSE):
                s = tf.compat.v1.reshape(s, [-1, 1])

                # embedding
                cell_initializer = tf.compat.v1.initializers.glorot_normal()

                h = tf.compat.v1.layers.dense(
                        s,
                        self.args.dis_dim,
                        use_bias=False,
                        kernel_initializer=cell_initializer,
                        name="layer_service_time")
                h = tf.compat.v1.nn.leaky_relu(h)  # [B*T, D_GAN]

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

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

        return h  # [B*T, 1]

    def _gradient_penalty_term(self, s_real, s_gen):

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

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

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

            return tf.compat.v1.abs(grad_penalty)

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

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

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

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

            return dis_loss

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

                gen_loss_1 = -self._logit_fake

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

                self.gen_loss_1 = tf.compat.v1.reduce_sum(gen_loss_1)
                self.gen_loss_2 = tf.constant(0.)
                self.gen_loss_3 = tf.constant(0.)

            return gen_loss

    def _predict_mark(self):
        """
        Args:
            output:
            covariates:
        """
        shape = tf.compat.v1.shape(self.input_data)
        batch_size = shape[0]
        seq_length = shape[1]

        next_mark = self._generator()
        next_mark = tf.compat.v1.reshape(next_mark, (batch_size, seq_length))

        return next_mark

    def _get_predicted_marks(self):
        shape = tf.compat.v1.shape(self.input_data)
        next_mark = self._generator()
        next_mark = tf.compat.v1.reshape(next_mark, (shape[0], shape[1]))
        next_mark = tf.compat.v1.split(next_mark, self.args.batch_size, axis=0)
        next_mark = tf.compat.v1.stack(next_mark)
        next_mark = tf.compat.v1.transpose(next_mark, [0, 2, 1])  # [B, T, n_mc_samples]

        return next_mark

    def _build_one_step_rmse(self):
        marks_real = self.input_data[:, :, -1]
        marks_predicted = self._predict_mark()
        marks_predicted = tf.compat.v1.multiply(self.mask, marks_predicted)

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


class RASBetaStaticCritic(DeepServiceTime):

    def __init__(self, args):
        """
        Args:
            args:
        """
        self.model = "RASBetaStaticCritic"
        super(RASBetaStaticCritic, self).__init__(args)
        self._build_service_distribution()
        self.sample_dist = self._get_predicted_marks()
        self.saver = tf.compat.v1.train.Saver(max_to_keep=100)

    def _build_recurrent(self):
        pass

    def _build_service_distribution(self):
        self._build_recurrent_adversarial()

    def _build_recurrent_adversarial(self):
        h = self.hidden_states_arrivals
        covariates = self.input_data[:, :, :-1]

        s_real = tf.compat.v1.reshape(self.input_data[:, :, -1],
                            [-1, 1])  # observed service time [B*T, 1]
        with tf.compat.v1.name_scope("build_generator-discriminator"):
            # Indices censored-uncensored data:
            s_real_sq = tf.compat.v1.squeeze(s_real)
            self._index_uncensored = tf.compat.v1.where(
                    tf.compat.v1.greater(s_real_sq, 0), tf.compat.v1.ones_like(s_real_sq), tf.compat.v1.zeros_like(s_real_sq))
            self._index_uncensored = tf.compat.v1.cast(self._index_uncensored, tf.compat.v1.int32)
            self._index_censored = tf.compat.v1.where(tf.compat.v1.less(s_real_sq, 0), tf.compat.v1.zeros_like(s_real_sq), tf.compat.v1.ones_like(s_real_sq))
            self._index_censored = tf.compat.v1.cast(self._index_censored, tf.compat.v1.int32)

            # Generator output
            self._s_gen, self.next_state = self._generator(
                    h, covariates, self.state)  # [B*T, 1], [B, T, H]

            # Discriminator output
            self._logit_real = self._discriminator(h, covariates, s_real)  # [B*T, 1]

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

            # Gradient penalty term:
            self._grad_penalty = self._gradient_penalty_term(s_real, self._s_gen, h, covariates)

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

        with tf.compat.v1.variable_scope("GAN/generator", reuse=tf.compat.v1.AUTO_REUSE):
            with tf.compat.v1.variable_scope("hidden_layers", reuse=tf.compat.v1.AUTO_REUSE):
                shape = tf.compat.v1.shape(h)
                B = shape[0]
                T = shape[1]

                cell_initializer = tf.compat.v1.initializers.glorot_normal()

                # embedding
                # sampling prior normal distr.
                z = tf.compat.v1.random_normal([B * T, self.args.gen_dim], 0.0, 1.0, dtype=tf.compat.v1.float64)
                z = tf.compat.v1.layers.dense(
                        z,
                        self.args.gen_dim,
                        use_bias=False,
                        kernel_initializer=cell_initializer,
                        name="layer_noise_0")

                if self.args.arrivals_hidden_state_bool:
                    h = tf.compat.v1.reshape(
                            h, [-1, self.args.hidden_units])

                    h = tf.compat.v1.layers.dense(
                            h,
                            self.args.gen_dim,
                            kernel_initializer=cell_initializer,
                            name="layer_hidden_0")
                    z = z + h

                if self.args.covariates_bool:
                    covariates = tf.compat.v1.reshape(covariates, [-1, self.args.input_dim - 1])
                    covariates = tf.compat.v1.layers.dense(
                            covariates,
                            self.args.gen_dim,
                            kernel_initializer=cell_initializer,
                            name="layer_emb_cov")
                    z = z + covariates

                # hidden state
                h = tf.compat.v1.nn.leaky_relu(z)  # [B*T, D_GAN]
                h = tf.compat.v1.reshape(h, [-1, self.args.num_steps, self.args.gen_dim])

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

                h = tf.compat.v1.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.compat.v1.layers.dense(
                                h,
                                self.args.gen_dim,
                                kernel_initializer=cell_initializer,
                                name="layer_hidden_{}".format(i + 1))
                        if self.args.noisy_layers:
                            z = tf.compat.v1.random_normal(
                                    tf.compat.v1.shape(h), 0.0, 1.0, dtype=tf.compat.v1.float64)
                            z = tf.compat.v1.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.compat.v1.nn.dropout(h, self.dropout_prob)
                        if self.args.batch_normalization:
                            h = tf.compat.v1.layers.batch_normalization(
                                    h, training=self.training_mode)
                        h = tf.compat.v1.nn.leaky_relu(h)  # [B*T, D_GAN]

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

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

    def _discriminator(self, h, covariates, s):
        """
        GAN discriminator
        h (arrival hidden state) [B, T, H]
        covariates [B, T, D-3]
        s (service time) [B*T, 1]
        """
        with tf.compat.v1.variable_scope("GAN/discriminator", reuse=tf.compat.v1.AUTO_REUSE):
            with tf.compat.v1.variable_scope("hidden_layers", reuse=tf.compat.v1.AUTO_REUSE):
                cell_initializer = tf.compat.v1.initializers.glorot_normal()

                # embedding
                s = tf.compat.v1.layers.dense(
                        s,
                        self.args.dis_dim,
                        kernel_initializer=cell_initializer,
                        name="layer_service_time")

                if self.args.arrivals_hidden_state_bool:
                    h = tf.compat.v1.reshape(h, [-1, self.args.hidden_units])
                    h = tf.compat.v1.layers.dense(
                            h,
                            self.args.dis_dim,
                            use_bias=False,
                            kernel_initializer=cell_initializer,
                            name="layer_hidden_0")
                    s = s + h

                if self.args.covariates_bool:
                    covariates = tf.compat.v1.reshape(covariates, [-1, self.args.input_dim - 1])
                    covariates = tf.compat.v1.layers.dense(
                            covariates,
                            self.args.dis_dim,
                            use_bias=False,
                            kernel_initializer=cell_initializer,
                            name="layer_emb_cov")
                    s = s + covariates

                # hidden state
                h = tf.compat.v1.nn.leaky_relu(s)  # [B*T, D_GAN]

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

            with tf.compat.v1.variable_scope(
                    "discriminator_output", reuse=tf.compat.v1.AUTO_REUSE):
                h = tf.compat.v1.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.compat.v1.name_scope("gradient_penalty_term"):
            eta = tf.compat.v1.random_uniform(
                    tf.compat.v1.shape(s_real), minval=0.0, maxval=0.1, dtype=tf.compat.v1.float64)

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

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

            return tf.compat.v1.abs(grad_penalty)

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

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

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

                dis_loss = tf.compat.v1.reshape(dis_loss, [batch_size, seq_length])
                dis_loss = tf.compat.v1.multiply(dis_loss, self.mask)
                dis_loss = tf.compat.v1.reshape(dis_loss, [-1, 1])

                dis_loss = tf.compat.v1.dynamic_partition(tf.compat.v1.squeeze(dis_loss), self._index_uncensored, 2)[1]

                dis_loss = tf.compat.v1.reduce_sum(dis_loss)

            return dis_loss

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

                # standard generator loss component of GAN
                gen_loss_1 = -self._logit_fake

                gen_loss_1 = tf.compat.v1.reshape(gen_loss_1, [batch_size, seq_length])
                gen_loss_1 = tf.compat.v1.multiply(gen_loss_1, self.mask)
                gen_loss_1 = tf.compat.v1.reshape(gen_loss_1, [-1, 1])

                _, gen_loss_1 = tf.compat.v1.dynamic_partition(tf.compat.v1.squeeze(gen_loss_1),
                                                     self._index_uncensored,
                                                     2)

                s_gen = self._s_gen

                # regulariser for censored events

                gen_loss_2 = tf.compat.v1.maximum(
                        tf.compat.v1.cast(0.0, tf.compat.v1.float64), self.T - s_gen)
                gen_loss_2 = tf.compat.v1.reshape(gen_loss_2, [batch_size, seq_length])
                gen_loss_2 = tf.compat.v1.multiply(gen_loss_2, self.mask)
                gen_loss_2 = tf.compat.v1.reshape(gen_loss_2, [-1, 1])

                _, gen_loss_2 = tf.compat.v1.dynamic_partition(tf.compat.v1.squeeze(gen_loss_2),
                                                     self._index_censored,
                                                     2)

                # regulariser for uncensored events

                s_real = tf.compat.v1.reshape(self.input_data[:, :, -1], [-1, 1])
                gen_loss_3 = tf.compat.v1.abs(s_real - s_gen)
                gen_loss_3 = tf.compat.v1.reshape(gen_loss_3, [batch_size, seq_length])
                gen_loss_3 = tf.compat.v1.multiply(gen_loss_3, self.mask)
                gen_loss_3 = tf.compat.v1.reshape(gen_loss_3, [-1, 1])

                _, gen_loss_3 = tf.compat.v1.dynamic_partition(tf.compat.v1.squeeze(gen_loss_3),
                                                     self._index_uncensored,
                                                     2)

                gen_loss = tf.compat.v1.reduce_sum(gen_loss_1) + self.beta_censoring * tf.compat.v1.reduce_sum(gen_loss_2) \
                           + tf.compat.v1.reduce_sum(gen_loss_3)

                self.gen_loss_1 = tf.compat.v1.reduce_sum(gen_loss_1)
                self.gen_loss_2 = tf.compat.v1.reduce_sum(gen_loss_2)
                self.gen_loss_3 = tf.compat.v1.reduce_sum(gen_loss_3)

            return gen_loss

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

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

        return next_mark, next_state

    def _get_predicted_marks(self):

        covariates = self.input_data[:, :, :-1]
        shape = tf.compat.v1.shape(self.hidden_states_arrivals)

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

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

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

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

        return next_mark

    def _build_one_step_rmse(self):

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

        marks_predicted, _ = self._predict_mark(
                self.hidden_states_arrivals, covariates)

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

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

        # Uncensored error:
        uncensored_error = tf.compat.v1.abs(marks_predicted - marks_real)

        # Censored error:
        censored_error = tf.compat.v1.maximum(tf.compat.v1.zeros_like(
                marks_predicted), self.T - marks_predicted)

        error = tf.compat.v1.where(tf.compat.v1.less(marks_real, 0),
                         censored_error, uncensored_error)

        self.one_step_rmse = error
