# Path: utils.py
# def get_seq_length(sequence):
#     used = tf.sign(tf.reduce_max(tf.abs(sequence), reduction_indices=2))
#     length = tf.reduce_sum(used, reduction_indices=1)
#     length = tf.cast(length, tf.int32)
#     return length
# 
# def _add_gradient_noise(t, stddev=1e-3, name=None):
#     """Adds gradient noise as described in http://arxiv.org/abs/1511.06807
#     The input Tensor `t` should be a gradient.
#     The output will be `t` + gaussian noise.
#     0.001 was said to be a good fixed value for memory networks."""
#     with tf.op_scope([t, stddev], name, "add_gradient_noise") as name:
#         t = tf.convert_to_tensor(t, name="t")
#         gn = tf.random_normal(tf.shape(t), stddev=stddev)
#         return tf.add(t, gn, name=name)
# 
# def _position_encoding(sentence_size, embedding_size):
#     """Position encoding described in section 4.1 in "End to End Memory Networks" (http://arxiv.org/pdf/1503.08895v5.pdf)"""
#     encoding = np.ones((embedding_size, sentence_size), dtype=np.float32)
#     ls = sentence_size + 1
#     le = embedding_size + 1
#     for i in range(1, le):
#         for j in range(1, ls):
#             encoding[i - 1, j - 1] = (i - (le - 1) / 2) * (j - (ls - 1) / 2)
#     encoding = 1 + 4 * encoding / embedding_size / sentence_size
#     return np.transpose(encoding)
# 
#     # TODO fix positional encoding so that it varies according to sentence lengths
# 
# def _xavier_weight_init():
#     """Xavier initializer for all variables except embeddings as desribed in [1]"""
# 
#     def _xavier_initializer(shape, **kwargs):
#         eps = np.sqrt(6) / np.sqrt(np.sum(shape))
#         out = tf.random_uniform(shape, minval=-eps, maxval=eps)
#         return out
# 
#     return _xavier_initializer
# 
# def _last_relevant(output, length):
#     """Finds the output at the end of each input"""
#     batch_size = int(output.get_shape()[0])
#     max_length = int(output.get_shape()[1])
#     out_size = int(output.get_shape()[2])
#     index = tf.range(0, batch_size) * max_length + (length - 1)
#     flat = tf.reshape(output, [-1, out_size])
#     relevant = tf.gather(flat, index)
#     return relevant
# 
# def batch_norm(x, is_training):
#     """ Batch normalization.
#     :param x: Tensor
#     :param is_training: boolean tf.Variable, true indicates training phase
#     :return: batch-normalized tensor
#     """
#     with tf.variable_scope('BatchNorm'):
#         # calculate dimensions (from tf.contrib.layers.batch_norm)
#         inputs_shape = x.get_shape()
#         axis = list(range(len(inputs_shape) - 1))
#         param_shape = inputs_shape[-1:]
# 
#         beta = tf.get_variable('beta', param_shape, initializer=tf.constant_initializer(0.))
#         gamma = tf.get_variable('gamma', param_shape, initializer=tf.constant_initializer(1.))
#         batch_mean, batch_var = tf.nn.moments(x, axis)
#         ema = tf.train.ExponentialMovingAverage(decay=0.5)
# 
#         def mean_var_with_update():
#             ema_apply_op = ema.apply([batch_mean, batch_var])
#             with tf.control_dependencies([ema_apply_op]):
#                 return tf.identity(batch_mean), tf.identity(batch_var)
# 
#         mean, var = tf.cond(is_training,
#                             mean_var_with_update,
#                             lambda: (ema.average(batch_mean), ema.average(batch_var)))
#         normed = tf.nn.batch_normalization(x, mean, var, beta, gamma, 1e-3)
#     return normed


import tensorflow as tf
from tensorflow.python.ops import  rnn, rnn_cell, seq2seq
from utils import get_seq_length, _add_gradient_noise, _position_encoding, _xavier_weight_init, _last_relevant, batch_norm
            return rnn_output

    #analogous to inference_step
    def generate_episode(self, batch_mask, prob_compare, prob, counter, episode, fact_vecs, acc_states, counter_int, weight_container, bias_container):    
        """Generate episode by applying attention to current fact vectors through a modified GRU"""
        fact_vecs_t = tf.unpack(tf.transpose(fact_vecs, perm=[1,0,2]))

        '''TRY REPLACING acc_states WITH episode AND SEE WHICH WORKS BETTER'''
        attentions = [tf.squeeze(self.get_attention(acc_states, fv), squeeze_dims=[1]) for fv in fact_vecs_t]

        attentions = tf.transpose(tf.pack(attentions))

        softs = tf.nn.softmax(attentions)
        softs = tf.split(1, self.max_input_len, softs)
        
        gru_outputs = []

        # set initial state to zero
        h = tf.zeros((self.batch_size, self.hidden_size))

        # use attention gru
        for i, fv in enumerate(fact_vecs_t):
            h = self._attention_GRU_step(fv, h, softs[i])
            gru_outputs.append(h)

        # extract gru outputs at proper index according to input_lens
        gru_outputs = tf.pack(gru_outputs)
        gru_outputs = tf.transpose(gru_outputs, perm=[1,0,2])

        #analogous to output, new_state = self.inference_cell(input,state)

# Path: utils.py
# def get_seq_length(sequence):
#     used = tf.sign(tf.reduce_max(tf.abs(sequence), reduction_indices=2))
#     length = tf.reduce_sum(used, reduction_indices=1)
#     length = tf.cast(length, tf.int32)
#     return length
# 
# def _add_gradient_noise(t, stddev=1e-3, name=None):
#     """Adds gradient noise as described in http://arxiv.org/abs/1511.06807
#     The input Tensor `t` should be a gradient.
#     The output will be `t` + gaussian noise.
#     0.001 was said to be a good fixed value for memory networks."""
#     with tf.op_scope([t, stddev], name, "add_gradient_noise") as name:
#         t = tf.convert_to_tensor(t, name="t")
#         gn = tf.random_normal(tf.shape(t), stddev=stddev)
#         return tf.add(t, gn, name=name)
# 
# def _position_encoding(sentence_size, embedding_size):
#     """Position encoding described in section 4.1 in "End to End Memory Networks" (http://arxiv.org/pdf/1503.08895v5.pdf)"""
#     encoding = np.ones((embedding_size, sentence_size), dtype=np.float32)
#     ls = sentence_size + 1
#     le = embedding_size + 1
#     for i in range(1, le):
#         for j in range(1, ls):
#             encoding[i - 1, j - 1] = (i - (le - 1) / 2) * (j - (ls - 1) / 2)
#     encoding = 1 + 4 * encoding / embedding_size / sentence_size
#     return np.transpose(encoding)
# 
#     # TODO fix positional encoding so that it varies according to sentence lengths
# 
# def _xavier_weight_init():
#     """Xavier initializer for all variables except embeddings as desribed in [1]"""
# 
#     def _xavier_initializer(shape, **kwargs):
#         eps = np.sqrt(6) / np.sqrt(np.sum(shape))
#         out = tf.random_uniform(shape, minval=-eps, maxval=eps)
#         return out
# 
#     return _xavier_initializer
# 
# def _last_relevant(output, length):
#     """Finds the output at the end of each input"""
#     batch_size = int(output.get_shape()[0])
#     max_length = int(output.get_shape()[1])
#     out_size = int(output.get_shape()[2])
#     index = tf.range(0, batch_size) * max_length + (length - 1)
#     flat = tf.reshape(output, [-1, out_size])
#     relevant = tf.gather(flat, index)
#     return relevant
# 
# def batch_norm(x, is_training):
#     """ Batch normalization.
#     :param x: Tensor
#     :param is_training: boolean tf.Variable, true indicates training phase
#     :return: batch-normalized tensor
#     """
#     with tf.variable_scope('BatchNorm'):
#         # calculate dimensions (from tf.contrib.layers.batch_norm)
#         inputs_shape = x.get_shape()
#         axis = list(range(len(inputs_shape) - 1))
#         param_shape = inputs_shape[-1:]
# 
#         beta = tf.get_variable('beta', param_shape, initializer=tf.constant_initializer(0.))
#         gamma = tf.get_variable('gamma', param_shape, initializer=tf.constant_initializer(1.))
#         batch_mean, batch_var = tf.nn.moments(x, axis)
#         ema = tf.train.ExponentialMovingAverage(decay=0.5)
# 
#         def mean_var_with_update():
#             ema_apply_op = ema.apply([batch_mean, batch_var])
#             with tf.control_dependencies([ema_apply_op]):
#                 return tf.identity(batch_mean), tf.identity(batch_var)
# 
#         mean, var = tf.cond(is_training,
#                             mean_var_with_update,
#                             lambda: (ema.average(batch_mean), ema.average(batch_var)))
#         normed = tf.nn.batch_normalization(x, mean, var, beta, gamma, 1e-3)
#     return normed


# Path: adaptive_attention.py
import tensorflow as tf
from tensorflow.python.ops import  rnn, rnn_cell, seq2seq
from utils import get_seq_length, _add_gradient_noise, _position_encoding, _xavier_weight_init, _last_relevant, batch_norm

            return rnn_output

    #analogous to inference_step
    def generate_episode(self, batch_mask, prob_compare, prob, counter, episode, fact_vecs, acc_states, counter_int, weight_container, bias_container):    
        """Generate episode by applying attention to current fact vectors through a modified GRU"""
        fact_vecs_t = tf.unpack(tf.transpose(fact_vecs, perm=[1,0,2]))

        '''TRY REPLACING acc_states WITH episode AND SEE WHICH WORKS BETTER'''
        attentions = [tf.squeeze(self.get_attention(acc_states, fv), squeeze_dims=[1]) for fv in fact_vecs_t]

        attentions = tf.transpose(tf.pack(attentions))

        softs = tf.nn.softmax(attentions)
        softs = tf.split(1, self.max_input_len, softs)
        
        gru_outputs = []

        # set initial state to zero
        h = tf.zeros((self.batch_size, self.hidden_size))

        # use attention gru
        for i, fv in enumerate(fact_vecs_t):
            h = self._attention_GRU_step(fv, h, softs[i])
            gru_outputs.append(h)

        # extract gru outputs at proper index according to input_lens
        gru_outputs = tf.pack(gru_outputs)
        gru_outputs = tf.transpose(gru_outputs, perm=[1,0,2])

        #analogous to output, new_state = self.inference_cell(input,state)

        episode = _last_relevant(gru_outputs, self.input_len_placeholder)