'''
Created on Mar 23, 2018

@author: anon
'''
from __future__ import absolute_import, division, print_function

import collections

from tensorflow.contrib.rnn.python.ops import core_rnn_cell
from tensorflow.python.framework import constant_op, dtypes
from tensorflow.python.framework.ops import convert_to_tensor
from tensorflow.python.ops import array_ops, gen_math_ops, random_ops, init_ops, control_flow_ops, gen_nn_ops, nn_ops, math_ops, rnn_cell_impl, variable_scope as vs
from tensorflow.python.ops.script_ops import py_func

from tensorflow.contrib.rnn import LSTMCell

import numpy as np
import tensorflow as tf

_Linear = core_rnn_cell._Linear

_CLSTMStateTuple = collections.namedtuple(
    "CLSTMStateTuple", ("c", "c_", "c_f", "h", "o", "decay", "t"))

sigmoid = math_ops.sigmoid
softplus = gen_nn_ops.softplus
add = math_ops.add
substract = math_ops.subtract
multiply = math_ops.multiply
exp = math_ops.exp
cond = control_flow_ops.cond
less_equal = gen_math_ops.less_equal


class CLSTMStateTuple(_CLSTMStateTuple):
    """Tuple used by LSTM Cells for `state_size`, `zero_state`, and output state.

    Stores two elements: `(c, c_, h, decay, t)`, in that order. Where `c` is the hidden state,
    `c_` is the base intensity, `decay` is the decaying rate, `h` is the output and
    `t` is the previous input.

    """
    __slots__ = ()

    @property
    def dtype(self):
        (c, c_, c_f, h, o, decay, t) = self
        if (c.dtype != h.dtype) or (c_.dtype != h.dtype):
            raise TypeError("Inconsistent internal state: %s vs %s vs %s" % 
                            (str(c.dtype), str(c_.dtype), str(h.dtype)))
        return c.dtype


class CLSTM(rnn_cell_impl.RNNCell):

    def __init__(self, num_units, activation=None, reuse=None, name=None):
        """Initialize the basic LSTM cell.

        Args:
          num_units: int, The number of units in the LSTM cell.
          activation: Activation function of the inner states.  Default: `tanh`.
          reuse: (optional) Python boolean describing whether to reuse variables
            in an existing scope.  If not `True`, and the existing scope already has
            the given variables, an error is raised.
          name: String, the name of the layer. Layers with the same name will
            share weights, but to avoid mistakes we require reuse=True in such
            cases.

          When restoring from CudnnLSTM-trained checkpoints, must use
          `CudnnCompatibleLSTMCell` instead.
        """
        super(CLSTM, self).__init__(_reuse=reuse, name=name)

        self._num_units = num_units
        self._activation = activation or math_ops.sigmoid

    @property
    def state_size(self):
        return CLSTMStateTuple(self._num_units, self._num_units,
                               self._num_units, self._num_units,
                               self._num_units, self._num_units, 1)

    @property
    def output_size(self):
        return tuple((self._num_units, self._num_units, self._num_units,
                      self._num_units, self._num_units))

    def new_output(self, c, c_, o, decay, t_prev, t):
        c_out = add(c_,
                    multiply(
                        substract(c, c_),
                        exp(-multiply(decay, substract(t, t_prev)))))
        h_out = multiply(o, substract(2 * sigmoid(2 * c_out), 1))
        return h_out, c_out

    def continuous_mode(self, inputs, state):
        c, c_, _, _, o, decay, t_prev = state
        t, _ = inputs
        h_out, _ = self.new_output(c, c_, o, decay, t_prev, t)
        #         h_out = Print(h_out, [h_out, "Continuous Mode"])
        return tuple((h_out, c, c_, decay, o)), state

    def discrete_mode(self, inputs, state):
        _, c_, c_f, h, _, _, t_prev = state
        t, k = inputs
        one = constant_op.constant(1, dtype=dtypes.int32)
        with vs.variable_scope('gates'):
            linear1 = _Linear([k, h], 6 * self._num_units, True)
            gate_inputs = linear1([k, h])
            gate_inputs = sigmoid(gate_inputs)
            i, i_, f, f_, z, o = array_ops.split(
                value=gate_inputs, num_or_size_splits=6, axis=one)
            z = substract(multiply(2., z), 1.)

        with vs.variable_scope('decay'):
            linear2 = _Linear([k, h], self._num_units, True)
            gate_inputs = linear2([k, h])
            decay = softplus(gate_inputs)
        c = add(multiply(f, c_f), multiply(i, z))
        c_ = add(multiply(f_, c_), multiply(i_, z))
        h_out, c_f = self.new_output(c, c_, o, decay, t_prev, t)
        new_state = CLSTMStateTuple(c, c_, c_f, h_out, o, decay, t)
        #         h_out = Print(h_out, ["Discrete Mode"])
        return tuple((h_out, c, c_, decay, o)), new_state

    def call(self, inputs, state):
        t, k = inputs
        # we choose the continous mode when all the entries of the event type (K) vector are 0.
        cond_expr = less_equal(math_ops.reduce_sum(k), 0)
        new_output, new_state = cond(cond_expr,
                                     lambda: self.continuous_mode(
                                         inputs, state),
                                     lambda: self.discrete_mode(inputs, state))

        return new_output, new_state


class EchoCell(rnn_cell_impl.RNNCell):

    def __init__(self,
                 num_units,
                 spectral_radius=.7,
                 reservoir_density=.1,
                 leaking_rate=1.,
                 noise=.001,
                 input_initializer=init_ops.random_normal_initializer(),
                 reservoir_initializer=init_ops.random_normal_initializer(),
                 feedback_initializer=init_ops.random_normal_initializer(),
                 bias_initializer=init_ops.random_normal_initializer(),
                 teacher_forcing=False,
                 output_dim=None,
                 activation=None,
                 reuse=None,
                 name=None,
                 state_is_tuple=False,
                 batch_size=None):

        super(EchoCell, self).__init__(_reuse=reuse, name=name)

        def _reservoir_initializer(shape, dtype, partition_info=None):
            W = reservoir_initializer(shape, dtype=dtype)

            connectivity_mask = math_ops.cast(
                math_ops.less_equal(
                    random_ops.random_uniform(shape), reservoir_density),
                dtype)

            W = math_ops.multiply(W, connectivity_mask)

            eigvals = py_func(np_eigenvals, [W], dtypes.complex64)
            rhoW = math_ops.reduce_max(math_ops.abs(eigvals))
            W *= self._spectral_radius / rhoW
            return W

        self._num_units = num_units
        self._output_dim = output_dim
        self._teacher_forcing = teacher_forcing
        self._spectral_radius = spectral_radius
        self._reservoir_density = reservoir_density
        self._leaking_rate = leaking_rate
        self._noise = noise
        self._input_initializer = input_initializer
        self._reservoir_initializer = _reservoir_initializer
        self._bias_initializer = bias_initializer
        self._feedback_initializer = feedback_initializer
        self._activation = activation or math_ops.tanh
        self._state_is_tuple = state_is_tuple

    @property
    def state_size(self):
        if self._state_is_tuple:
            return (self._num_units, self._num_units)
        else:
            return self._num_units

    @property
    def output_size(self):
        return self._num_units

    def call(self, inputs, state):
        if self._state_is_tuple:
            state_tensor = state[0]
        else:
            state_tensor = state
        inputs = convert_to_tensor(inputs)
        input_depth = inputs.shape[1]
        dtype = inputs.dtype
        self._build(input_depth, dtype)
        res = self._linear_multiplicatio_activate(inputs, state_tensor)
        output = add(
            multiply(substract(1., self._leaking_rate), state_tensor),
            multiply(self._leaking_rate, res))
        # h, [h, c]
        if self._state_is_tuple:
            state = (output, output)
            return output, state
        else:
            return output, output

    def _build(self, input_depth, dtype):
        if self._teacher_forcing:
            self._W_feedback = vs.get_variable(
                'feedback_weight', [self._output_dim, self._num_units],
                dtype,
                self._feedback_initializer,
                trainable=False)
            input_depth -= self._output_dim

        self._W_in = vs.get_variable(
            'input_weight', [input_depth, self._num_units],
            dtype,
            self._input_initializer,
            trainable=False)
        self._W = vs.get_variable(
            'reservoir', [self._num_units, self._num_units],
            dtype,
            self._reservoir_initializer,
            trainable=False)
        self._W_bias = vs.get_variable(
            'bias_weight', [self._num_units],
            dtype,
            self._bias_initializer,
            trainable=False)

    def _linear_multiplicatio_activate(self, inputs, state):
        if self._state_is_tuple:
            state = state
        one = constant_op.constant(1, dtypes.int32)
        zero = constant_op.constant(0, dtypes.int32)
        in_matrix = array_ops.concat([inputs, state], axis=one)
        if self._teacher_forcing:
            w_matrix = array_ops.concat(
                [self._W_in, self._W_feedback, self._W], axis=zero)
        else:
            w_matrix = array_ops.concat([self._W_in, self._W], axis=zero)
        res = math_ops.matmul(in_matrix, w_matrix)
        res = nn_ops.bias_add(res, self._W_bias)

        res = add(res,
                  random_ops.random_uniform(tf.shape(res), -self._noise,
                                            self._noise))
        res = self._activation(res)
        return res


class VariationalRNNCell(tf.contrib.rnn.RNNCell):

    # def __init__(self, x_dim, h_dim, z_dim=100):
    def __init__(self, h_dim, z_dim, NL, args, x_dim=2):

        self.n_h = h_dim

        self.n_x = x_dim
        self.n_z = z_dim
        # self.n_x_1 = x_dim
        # self.n_z_1 = z_dim
        self.n_enc_hidden = z_dim
        self.n_dec_hidden = x_dim
        self.n_prior_hidden = z_dim

        self.Wz = generate_initialized_variable([z_dim, h_dim], 'Wz')
        self.Wx = generate_initialized_variable([x_dim, h_dim], 'Wx')

        # self.lstm = tf.contrib.rnn.LSTMCell(self.n_h, state_is_tuple=True)
        if args.esn:
            self.rnn = EchoCell(self.n_h)
        else:
            self.rnn = tf.contrib.rnn.BasicRNNCell(self.n_h, NL)

    @property
    def state_size(self):
        return self.n_h

    @property
    def output_size(self):
        # return self.n_h
        return (self.n_h, self.n_z, self.n_z, self.n_z, self.n_x, self.n_x, self.n_z, self.n_z)

    def __call__(self, x, state):

        prior_mu, prior_sigma = self.prior(state)
        # # x_1 = self.phi_x(x)
        enc_mu, enc_sigma = self.encode(x, state)
        z = self.reparameterize(enc_mu, enc_sigma, x)
        # # z_1 = self.phi_z(z)
        dec_mu, dec_sigma = self.decode(z, state)
        Wz = tf.matmul(z, self.Wz)
        Wx = tf.matmul(x, self.Wx)
        output, state2 = self.rnn(Wz + Wx, state)
        # output, state2 = self.rnn(Wx, state)

        output = (output, z, enc_mu, enc_sigma, dec_mu,
                  dec_sigma, prior_mu, prior_sigma)

        return output, state2

        # state2 = self.lstm(x, state)
        # return state2

    def prior(self, h):
        with tf.variable_scope("Prior"):
            with tf.variable_scope("hidden"):
                prior_hidden = tf.nn.relu(linear(h, self.n_prior_hidden))
            with tf.variable_scope("mu"):
                prior_mu = linear(prior_hidden, self.n_z)
            with tf.variable_scope("sigma"):
                prior_sigma = tf.nn.softplus(linear(prior_hidden, self.n_z))
        return prior_mu, prior_sigma

    def phi_x(self, x):
        with tf.variable_scope("Phi_x"):
            x_1 = tf.nn.relu(linear(x, self.n_x_1))
        return x_1

    def phi_z(self, z):
        with tf.variable_scope("Phi_z"):
            z_1 = tf.nn.relu(linear(z, self.n_z_1))
        return z_1

    def encode(self, x, h):
        with tf.variable_scope("Encoder"):
            with tf.variable_scope("hidden"):
                enc_hidden = tf.nn.relu(linear(tf.concat(axis=1,
                                                         values=(x, h)), self.n_enc_hidden))
            with tf.variable_scope("mu"):
                enc_mu = linear(enc_hidden, self.n_z)
            with tf.variable_scope("sigma"):
                enc_sigma = tf.nn.softplus(linear(enc_hidden, self.n_z))
        return enc_mu, enc_sigma

    def reparameterize(self, enc_mu, enc_sigma, x):
        # TODO where to incorporate prior(z)
        eps = tf.random_normal(
            (x.get_shape().as_list()[0], self.n_z), 0.0, 1.0, dtype=tf.float32)
        z = tf.add(enc_mu, tf.multiply(enc_sigma, eps))
        return z

    def decode(self, z, h):
        with tf.variable_scope("Decoder"):
            with tf.variable_scope("hidden"):
                dec_hidden = tf.nn.relu(
                    linear(tf.concat(axis=1, values=(z, h)), self.n_dec_hidden))
            with tf.variable_scope("mu"):
                dec_mu = linear(dec_hidden, self.n_x)
            with tf.variable_scope("sigma"):
                dec_sigma = tf.nn.softplus(linear(dec_hidden, self.n_x))

        return dec_mu, dec_sigma


def generate_initialized_variable(shape, name, dtype=tf.float32):
    initialization = tf.random_normal(shape, stddev=.1, dtype=dtype)
    return tf.Variable(initialization, name=name)


def linear(input_, output_size, scope=None, stddev=0.02,
           bias_start=0.0, with_w=False):
    shape = input_.get_shape().as_list()
    with tf.variable_scope(scope or "Linear"):
        matrix = tf.get_variable("Matrix", [shape[1], output_size], tf.float32,
                                 tf.random_normal_initializer(stddev=stddev))
        bias = tf.get_variable("bias", [output_size],
                               initializer=tf.constant_initializer(bias_start))
        if with_w:
            return tf.matmul(input_, matrix) + bias, matrix, bias
        else:
            return tf.matmul(input_, matrix) + bias


def np_eigenvals(x):
    return np.linalg.eigvals(x).astype('complex64')


class GeneratorRecurrentServiceTime(tf.contrib.rnn.RNNCell):

    def __init__(self,
                 batch_size,
                 h_dim,
                 service_load_size,
                 n_layers,
                 noisy_layers,
                 dropout,
                 dropout_prob,
                 batch_normalization,
                 training):

        """
        Generator cell for recurrent service time model

        """
        self._batch_size = batch_size

        """ System Load parameters """

        """ Dimensions """

        self._h_dim = h_dim
        self._service_load_size = service_load_size

        """ Layers """

        self._n_layers = n_layers
        self._noisy_layers = noisy_layers

        """ Nonlinearity """

        self._nonlinearity = tf.nn.relu

        """ Weight & bias initializers """

        self._weight_initializer = tf.keras.initializers.he_normal(seed=None)

        """ Recurrent cell definition """

        with tf.variable_scope("generator/rnn", reuse=tf.AUTO_REUSE):
            self._rnn_cell = LSTMCell(self._h_dim, state_is_tuple=True)

        """ Dropout """

        self._dropout = dropout
        self._dropout_prob = dropout_prob

        if self._dropout:
            self._rnn_cell = tf.nn.rnn_cell.DropoutWrapper(self._rnn_cell, output_keep_prob=self._dropout_prob)

        """ Batch normalization """

        # Boolean variables defining batch normalization
        self._batch_normalization = batch_normalization  # python boolean
        self._training_mode = training  # tensor boolean

    @property
    def state_size(self):
        return self._h_dim, self._h_dim

    @property
    def output_size(self):
        """
        output: prediction of generator
        """
        return self._service_load_size

    def get_initial_state(self):
        """
        Initializes the cell state
        """
        with tf.name_scope("initial_state"):
            c = tf.zeros([self._batch_size, self._h_dim],
                         name="cell")
            h = tf.zeros([self._batch_size, self._h_dim],
                         name="hidden")

        initial_state = tf.nn.rnn_cell.LSTMStateTuple(c, h)

        return initial_state

    def get_state_placeholder(self):

        with tf.name_scope("initial_state_placeholder"):
            cp = tf.placeholder(tf.float64,
                                shape=[None, self._h_dim],
                                name="cell")
            hp = tf.placeholder(tf.float64,
                                shape=[None, self._h_dim],
                                name="hidden")

        initial_state_placeholder = tf.nn.rnn_cell.LSTMStateTuple(cp, hp)

        return initial_state_placeholder

    def __call__(self, input, state):
        """
        Generator for recurrent service time cell with recurrent GANs

        input shape: [B, D]

        state: (c, h)

        """
        with tf.variable_scope("generator_call", reuse=tf.AUTO_REUSE):
            # Prior noise [B, h]:
            z = tf.random_normal((self._batch_size, self._h_dim), 0.0, 1.0, dtype=tf.float64)

            # Embedding of noise and input:
            with tf.variable_scope("embedding", reuse=tf.AUTO_REUSE):
                z = tf.layers.dense(z, self._h_dim,
                                    kernel_initializer=self._weight_initializer,
                                    name="layer_hidden_z")

                input = tf.layers.dense(input, self._h_dim,
                                        kernel_initializer=self._weight_initializer,
                                        use_bias=False,
                                        name="layer_hidden_input")

                h = self._nonlinearity(z + input)

            # Recurrent cell:
            with tf.name_scope("generator_rnn"):
                _, state = self._rnn_cell(h, state)
            _, h = state

            if self._n_layers > 1:
                with tf.variable_scope("output_rnn_layers", reuse=tf.AUTO_REUSE):
                    for i in range(self._n_layers - 1):
                        h = tf.layers.dense(h, self._h_dim,
                                            kernel_initializer=self._weight_initializer,
                                            name="layer_hidden_{}".format(i + 1))
                        if self._noisy_layers:
                            z = tf.random_normal(tf.shape(h), 0.0, 1.0, dtype=tf.float32)
                            z = tf.layers.dense(z, self._h_dim, use_bias=False,
                                                kernel_initializer=self._weight_initializer,
                                                name="layer_noise_{}".format(i + 1))
                            h = h + z
                        if self._dropout:
                            h = tf.nn.dropout(h, self._dropout_prob)
                        if self._batch_normalization:
                            h = tf.layers.batch_normalization(h, training=self._training_mode)
                        h = self._nonlinearity(h)

            # Output:
            with tf.variable_scope("final_output", reuse=tf.AUTO_REUSE):
                h = tf.layers.dense(h, self._service_load_size, kernel_initializer=self._weight_initializer)
                h = tf.exp(h)  # [B, service_load_size]

        return h, state


class DiscriminatorRecurrentServiceTime(tf.contrib.rnn.RNNCell):

    def __init__(self,
                 batch_size,
                 h_dim,
                 n_layers,
                 dropout,
                 dropout_prob,
                 batch_normalization,
                 training):

        """
        Discriminator (critic) cell for recurrent service time model

        """
        self._batch_size = batch_size

        """ Dimensions """

        self._h_dim = h_dim

        """ Layers """

        self._n_layers = n_layers

        """ Nonlinearity """

        self._nonlinearity = tf.nn.leaky_relu

        """ Weight initializers """

        self._weight_initializer = tf.keras.initializers.he_normal(seed=None)

        """ Recurrent cell definition """

        with tf.variable_scope("discriminator/rnn", reuse=tf.AUTO_REUSE):
            self._rnn_cell = LSTMCell(self._h_dim, state_is_tuple=True)

        """ Dropout """

        self._dropout = dropout
        self._dropout_prob = dropout_prob

        if self._dropout:
            self._rnn_cell = tf.nn.rnn_cell.DropoutWrapper(self._rnn_cell, output_keep_prob=self._dropout_prob)

        """ Batch normalization """

        # Boolean variables defining batch normalization
        self._batch_normalization = batch_normalization  # python boolean
        self._training_mode = training  # tensor boolean

    @property
    def state_size(self):
        return self._h_dim, self._h_dim

    @property
    def output_size(self):
        """
        output: critic output
        """
        return 1

    def get_initial_state(self):
        """
        Initializes the cell state
        """
        with tf.name_scope("initial_state"):
            c = tf.zeros([self._batch_size, self._h_dim],
                         name="cell")
            h = tf.zeros([self._batch_size, self._h_dim],
                         name="hidden")

        initial_state = tf.nn.rnn_cell.LSTMStateTuple(c, h)

        return initial_state

    def get_state_placeholder(self):

        with tf.name_scope("initial_state_placeholder"):
            cp = tf.placeholder(tf.float64,
                                shape=[None, self._h_dim],
                                name="cell")
            hp = tf.placeholder(tf.float64,
                                shape=[None, self._h_dim],
                                name="hidden")

        initial_state_placeholder = tf.nn.rnn_cell.LSTMStateTuple(cp, hp)

        return initial_state_placeholder

    def __call__(self, input, state):
        """
        Discriminator for recurrent service time cell with recurrent GANs

        input shape: [B, D]

        state: (c, h)

        """
        with tf.variable_scope("discriminator_call", reuse=tf.AUTO_REUSE):

            # Embedding
            with tf.variable_scope("embedding", reuse=tf.AUTO_REUSE):
                input = tf.layers.dense(input, self._h_dim,
                                        kernel_initializer=self._weight_initializer,
                                        use_bias=True,
                                        name="layer_hidden_input")

                h = self._nonlinearity(input)

            # Recurrent cell:
            with tf.name_scope("discriminator_rnn"):
                _, state = self._rnn_cell(h, state)
            _, h = state

            if self._n_layers > 1:
                with tf.variable_scope("output_rnn_layers", reuse=tf.AUTO_REUSE):
                    for i in range(self._n_layers - 1):
                        h = tf.layers.dense(h, self._h_dim,
                                            kernel_initializer=self._weight_initializer,
                                            name="layer_hidden_{}".format(i + 1))
                        if self._dropout:
                            h = tf.nn.dropout(h, self._dropout_prob)
                        if self._batch_normalization:
                            h = tf.layers.batch_normalization(h, training=self._training_mode)
                        h = self._nonlinearity(h)

            with tf.variable_scope("final_output", reuse=tf.AUTO_REUSE):
                h = tf.layers.dense(h, 1, kernel_initializer=self._weight_initializer)

        return h, state
