from collections import namedtuple
import tensorflow as tf


class TimedNTM(object):
    """Timed Neural Turing Machine

    inspired by paper:
        https://arxiv.org/pdf/1410.5401.pdf
    implementation inspired by:
        https://github.com/carpedm20/NTM-tensorflow/blob/master/ntm_cell.py
    """

    def __init__(self, attn_shift_range, sparse_attention, name):
        """Construct the `TimedNTM`.

        Args:
            attn_shift_range: Python int.
                A time range within which to attend to the memory by location
            sparse_attention: Python bool.
                If `True` use sparsemax instead of softmax for probs
            name: Name to use when creating ops.
        """

        # interpolation gate
        self.name = 'timed_ntm_' + name

        self._inter_gate = tf.layers.Dense(
                units=1,
                activation=tf.sigmoid,
                name=self.name + '/inter_gate'
        )
        # if use sparsemax instead of softmax for probs
        self._sparse_attention = sparse_attention

        if sparse_attention:
            # sparsemax doesn't support inf
            self._inf = float(5000)
        else:
            self._inf = float('inf')

        # shift weighting if range is provided
        if attn_shift_range:
            self._shift_weight = tf.layers.Dense(
                    units=2 * attn_shift_range + 1,
                    activation=tf.nn.softmax,
                    name=self.name + '/shift_weight'
            )
        else:
            self._shift_weight = None

        # sharpening parameter
        self._gamma_sharp = tf.layers.Dense(
                units=1,
                activation=lambda a: tf.nn.softplus(a) + 1,
                bias_initializer=tf.constant_initializer(1),
                name=self.name + '/gamma_sharp'
        )

    def __call__(self, attn_inputs, scores, scores_state, mask):
        # apply exponential moving average with interpolation gate weight
        # to scores from previous time which are equal to probs at this point
        # different from original NTM where it is applied after softmax
        i_g = self._inter_gate(attn_inputs)

        # scores limited by time
        scores = tf.concat([i_g * scores[:, :-1] + (1 - i_g) * scores_state,
                            scores[:, -1:]], 1)
        next_scores_state = scores

        if mask is not None:
            # apply mask to scores
            if self._shift_weight is not None:
                # rearrange scores to make them continuous for convolution
                scores = tf.map_fn(self._rearrange_fn,
                                   [scores, mask], dtype=scores.dtype)
            else:
                scores = tf.where(mask > 0,
                                  scores, -self._inf * tf.ones_like(scores))

        # create probabilities for attention
        if self._sparse_attention:
            probs = tf.contrib.sparsemax.sparsemax(scores)
        else:
            probs = tf.nn.softmax(scores)

        if self._shift_weight is not None:
            s_w = self._shift_weight(attn_inputs)

            # we want to go back in time during convolution
            conv_probs = tf.reverse(probs, axis=[1])

            # preare probs for tf.nn.depthwise_conv2d
            # [in_width, in_channels=batch]
            conv_probs = tf.transpose(conv_probs, [1, 0])
            # [batch=1, in_height=1, in_width=time+1, in_channels=batch]
            conv_probs = conv_probs[tf.newaxis, tf.newaxis, :, :]

            # [filter_height=1, filter_width=2*attn_shift_range+1,
            #   in_channels=batch, channel_multiplier=1]
            conv_s_w = tf.transpose(s_w, [1, 0])
            conv_s_w = conv_s_w[tf.newaxis, :, :, tf.newaxis]

            # perform 1d convolution
            # [batch=1, out_height=1, out_width=time+1, out_channels=batch]
            conv_probs = tf.nn.depthwise_conv2d_native(conv_probs, conv_s_w,
                                                       [1, 1, 1, 1], 'SAME')
            conv_probs = conv_probs[0, 0, :, :]
            conv_probs = tf.transpose(conv_probs, [1, 0])

            probs = tf.reverse(conv_probs, axis=[1])

            if mask is not None:
                # arrange probs back to their original time order
                probs = tf.map_fn(self._arrange_back_fn,
                                  [probs, mask], dtype=probs.dtype)

        # sharpening
        g_sh = self._gamma_sharp(attn_inputs)

        powed_probs = tf.pow(probs, g_sh)
        probs = powed_probs / (
                tf.reduce_sum(powed_probs, 1, keepdims=True) + 1e-32)

        return probs, next_scores_state

    def _rearrange_fn(self, list_tensor_1d_mask_1d):
        """Rearranges tensor_1d to put all the values
            where mask_1d=1 to the right and
            where mask_1d=0 to the left and sets them to -infinity"""
        tensor_1d, mask_1d = list_tensor_1d_mask_1d

        partitioned_tensor = tf.dynamic_partition(tensor_1d,
                                                  mask_1d, 2)
        partitioned_tensor[0] = \
            -self._inf * tf.ones_like(partitioned_tensor[0])

        return tf.concat(partitioned_tensor, 0)

    @staticmethod
    def _arrange_back_fn(list_tensor_1d_mask_1d):
        """Arranges back tensor_1d to restore original order
            modified by `_rearrange_fn` according to mask_1d:
            - number of 0s in mask_1d values on the left are set to
              their corresponding places where mask_1d=0,
            - number of 1s in mask_1d values on the right are set to
              their corresponding places where mask_1d=1"""
        tensor_1d, mask_1d = list_tensor_1d_mask_1d

        mask_indices = tf.dynamic_partition(tf.range(tf.shape(tensor_1d)[0]),
                                            mask_1d, 2)

        mask_sum = tf.reduce_sum(mask_1d, axis=0)
        partitioned_tensor = [tf.zeros_like(tensor_1d[:-mask_sum]),
                              tensor_1d[-mask_sum:]]

        return tf.dynamic_stitch(mask_indices, partitioned_tensor)


def _compute_time_attention(attention_mechanism, attn_inputs, attention_state,
                            # time is added to calculate time attention
                            time, timed_ntm, time_mask, ignore_mask,
                            attention_layer):
    """Computes the attention and alignments limited by time
        for a given attention_mechanism.

        Modified helper method from tensorflow."""

    scores, _ = attention_mechanism(attn_inputs, state=attention_state)

    # take only scores from current and past times
    timed_scores = scores[:, :time + 1]
    timed_scores_state = attention_state[:, :time]

    # get mask for past times
    timed_time_mask = time_mask[:, :time]
    if ignore_mask is not None:
        timed_time_mask *= 1 - ignore_mask[:, :time]

    # set mask for current time to 1
    timed_time_mask = tf.concat([timed_time_mask,
                                 tf.ones_like(time_mask[:, :1])], 1)

    # pass these scores to NTM
    probs, next_scores_state = timed_ntm(attn_inputs, timed_scores,
                                         timed_scores_state,
                                         timed_time_mask)

    # concatenate probs with zeros to get new alignments
    zeros = tf.zeros_like(scores)
    # remove current time from attention
    alignments = tf.concat([probs[:, :-1], zeros[:, time:]], 1)

    # Reshape from [batch_size, memory_time] to [batch_size, 1, memory_time]
    expanded_alignments = tf.expand_dims(alignments, 1)

    # Context is the inner product of alignments and values along the
    # memory time dimension.
    # alignments shape is
    #   [batch_size, 1, memory_time]
    # attention_mechanism.values shape is
    #   [batch_size, memory_time, memory_size]
    # the batched matmul is over memory_time, so the output shape is
    #   [batch_size, 1, memory_size].
    # we then squeeze out the singleton dim.
    context = tf.matmul(expanded_alignments, attention_mechanism.values)
    context = tf.squeeze(context, [1])

    if attention_layer is not None:
        attention = attention_layer(tf.concat([attn_inputs, context], 1))
    else:
        attention = context

    # return current time to attention
    alignments = tf.concat([probs, zeros[:, time + 1:]], 1)
    next_attention_state = tf.concat([next_scores_state,
                                      zeros[:, time + 1:]], 1)
    return attention, alignments, next_attention_state


# noinspection PyProtectedMember
class TimeAttentionWrapperState(
            namedtuple("TimeAttentionWrapperState",
                       tf.contrib.seq2seq.AttentionWrapperState._fields +
                       ("all_time_masks", "all_cell_states"))):  # added
    """Modified  from tensorflow's tf.contrib.seq2seq.AttentionWrapperState
        see there for description of the parameters

    Additional fields:
        - `all_time_masks`: A mask applied to a memory
           that filters certain time steps
        - `all_cell_states`: All states of the wrapped `RNNCell`
           at all the previous time steps.
    """

    def clone(self, **kwargs):
        """Copied  from tensorflow's tf.contrib.seq2seq.AttentionWrapperState
            see there for description of the parameters"""

        def with_same_shape(old, new):
            """Check and set new tensor's shape."""
            if isinstance(old, tf.Tensor) and isinstance(new, tf.Tensor):
                return tf.contrib.framework.with_same_shape(old, new)
            return new

        return tf.contrib.framework.nest.map_structure(
                with_same_shape,
                self,
                super(TimeAttentionWrapperState, self)._replace(**kwargs)
        )


class TimeAttentionWrapper(tf.contrib.seq2seq.AttentionWrapper):
    """Custom AttentionWrapper that takes into account time
        when calculating attention.
        Attention is calculated before calling rnn cell.

        Modified from tensorflow's tf.contrib.seq2seq.AttentionWrapper.
    """

    def __init__(self, cell,
                 attention_mechanism,
                 sequence_len,
                 attn_shift_range=0,
                 sparse_attention=False,
                 attention_layer_size=None,
                 alignment_history=False,
                 rnn_and_attn_inputs_fn=None,
                 ignore_mask=None,
                 cell_input_fn=None,
                 index_of_attn_to_copy=None,
                 likelihood_fn=None,
                 tensor_not_to_copy=None,
                 output_attention=False,
                 initial_cell_state=None,
                 name=None,
                 attention_layer=None):
        """Construct the `TimeAttentionWrapper`.
            See the super class for the original arguments description.

        Additional args:
            sequence_len: Python integer.
                Maximum length of the sequence, used to create
                appropriate TensorArray for all cell states
                in TimeAttentionWrapperState
            attn_shift_range: Python integer (`0` by default).
                A time range within which to attend to the memory
                by location in Neural Turing Machine.
            sparse_attention: Python bool.
                A flag to use sparsemax (if `True`) instead of
                softmax (if `False`, default) for probabilities
            inputs_and_attn_inputs_fn: (optional) A `callable`.
                A function that creates inputs and attention inputs tensors.
            ignore_mask: (optional) Boolean Tensor.
                Determines which time steps to ignore in attention
            index_of_attn_to_copy: (optional) Python integer.
                An index of attention mechanism that picks
                which part of attention tensor to use for copying to output,
                the default is `None`, which turns off copying mechanism.
                Copy inspired by: https://arxiv.org/pdf/1603.06393.pdf
            likelihood_fn: (optional) A `callable`.
                A method to perform likelihood calculation to
                filter time step in copy mechanism.
                Returns a tuple of binary likelihood and likelihood
            tensor_not_to_copy: (optional) A Tensor.
                A tensor, which shouldn't be copied from previous time steps

        Modified args:
            output_attention: Python bool.  If `True`, the output at each
                time step is the concatenated cell outputs,
                attention values and additional values described in
                `additional_output_size()`, used in copy mechanism.
        """
        super(TimeAttentionWrapper, self).__init__(
                cell,
                attention_mechanism,
                attention_layer_size,
                alignment_history,
                cell_input_fn,
                output_attention,
                initial_cell_state,
                name,
                attention_layer
        )
        self._sequence_len = sequence_len

        if not isinstance(attn_shift_range, list):
            # attn_shift_range might not be a list
            attn_shift_range = [attn_shift_range]
        self._timed_ntms = [TimedNTM(attn_shift_range[0],
                                     sparse_attention,
                                     name='0')]
        if self._is_multi:
            # if there are several attention mechanisms,
            # create additional TimedNTMs for them
            if len(attn_shift_range) == 1:
                # original attn_shift_range might not be a list
                attn_shift_range *= len(attention_mechanism)
            elif len(attn_shift_range) != len(attention_mechanism):
                raise ValueError(
                    "If provided, `attn_shift_range` must contain exactly one "
                    "integer per attention_mechanism, saw: {} vs {}"
                    "".format(len(attn_shift_range), len(attention_mechanism))
                )
            for i in range(1, len(attention_mechanism)):
                self._timed_ntms.append(TimedNTM(attn_shift_range[i],
                                                 sparse_attention,
                                                 name=str(i)))

        if rnn_and_attn_inputs_fn is None:
            rnn_and_attn_inputs_fn = self._default_rnn_and_attn_inputs_fn
        else:
            if not callable(rnn_and_attn_inputs_fn):
                raise TypeError(
                    "`rnn_and_attn_inputs_fn` must be callable, saw type: {}"
                    "".format(type(rnn_and_attn_inputs_fn).__name__)
                )
        self._rnn_and_attn_inputs_fn = rnn_and_attn_inputs_fn

        if not isinstance(ignore_mask, list):
            self._ignore_mask = [tf.cast(ignore_mask, tf.int32)]
        else:
            self._ignore_mask = [tf.cast(i_m, tf.int32) for i_m in ignore_mask]

        self._index_of_attn_to_copy = index_of_attn_to_copy

        self._likelihood_fn = likelihood_fn
        self._tensor_not_to_copy = tensor_not_to_copy

    @staticmethod
    def _default_rnn_and_attn_inputs_fn(inputs, cell_state):
        if isinstance(cell_state, tf.contrib.rnn.LSTMStateTuple):
            return inputs, tf.concat([inputs, cell_state.h], -1)
        else:
            return inputs, tf.concat([inputs, cell_state], -1)

    @staticmethod
    def additional_output_size():
        """Number of additional outputs:

        likelihoods:
            attn_likelihood, state_likelihood
        debugging info:
            current_time_prob,
            bin_likelihood_not_to_copy, bin_likelihood_to_copy

        **Method should be static**
        """
        return 2 + 3

    @property
    def output_size(self):
        if self._output_attention:
            if self._index_of_attn_to_copy is not None:
                # output both raw rnn cell_output and
                # cell_output with copied attention
                # together with attention vector itself
                # and additional output
                return (2 * self._cell.output_size +
                        self._attention_layer_size +
                        self.additional_output_size())
            else:
                return self._cell.output_size + self._attention_layer_size
        else:
            return self._cell.output_size

    @property
    def state_size(self):
        """The `state_size` property of `TimeAttentionWrapper`.
        Returns:
            A `TimeAttentionWrapperState` tuple containing shapes
            used by this object.
        """

        # use AttentionWrapperState from superclass
        state_size = super(TimeAttentionWrapper, self).state_size

        all_cell_states = self._cell.state_size

        return TimeAttentionWrapperState(
                cell_state=state_size.cell_state,
                time=state_size.time,
                attention=state_size.attention,
                alignments=state_size.alignments,
                attention_state=state_size.attention_state,
                alignment_history=state_size.alignment_history,
                all_time_masks=self._sequence_len,
                all_cell_states=all_cell_states)

    def zero_state(self, batch_size, dtype):
        """Modified  from tensorflow's zero_state
            see there for description of the parameters"""

        # use AttentionWrapperState from superclass
        zero_state = super(TimeAttentionWrapper,
                           self).zero_state(batch_size, dtype)

        with tf.name_scope(type(self).__name__ + "ZeroState",
                           values=[batch_size]):
            # store time masks
            all_time_masks = tf.TensorArray(
                    tf.int32,
                    size=self._sequence_len + 1,
                    dynamic_size=False,
                    clear_after_read=False
            ).write(0, tf.zeros([batch_size, self.state_size.all_time_masks],
                                tf.int32))

            # store all cell states into a tensor array to allow
            # copy mechanism to go back in time
            if isinstance(self._cell.state_size,
                          tf.contrib.rnn.LSTMStateTuple):
                all_cell_states = tf.contrib.rnn.LSTMStateTuple(
                        tf.TensorArray(dtype, size=self._sequence_len + 1,
                                       dynamic_size=False,
                                       clear_after_read=False
                                       ).write(0, zero_state.cell_state.c),
                        tf.TensorArray(dtype, size=self._sequence_len + 1,
                                       dynamic_size=False,
                                       clear_after_read=False
                                       ).write(0, zero_state.cell_state.h)
                )
            else:
                all_cell_states = tf.TensorArray(
                        dtype, size=0,
                        dynamic_size=False,
                        clear_after_read=False
                ).write(0, zero_state.cell_state)

            return TimeAttentionWrapperState(
                    cell_state=zero_state.cell_state,
                    time=zero_state.time,
                    attention=zero_state.attention,
                    alignments=zero_state.alignments,
                    attention_state=zero_state.attention_state,
                    alignment_history=zero_state.alignment_history,
                    all_time_masks=all_time_masks,
                    all_cell_states=all_cell_states
            )

    def call(self, inputs, state):
        """Perform a step of attention-wrapped RNN.

        The order has changed:
        - Step 1: Calculate attention inputs based on the previous cell state
                  and current inputs
        - Step 2: Score the output with `attention_mechanism`.
        - Step 3: Calculate the alignments by passing the score through the
                  `normalizer` and limit them by time.
        - Step 4: Calculate the context vector as the inner product between the
                  alignments and the attention_mechanism's values (memory).
        - Step 5: Calculate the attention output by concatenating
                  the cell output and context through the attention layer
                  (a linear layer with `attention_layer_size` outputs).
        - Step 6: Mix the `inputs` and `attention` output via
                  `cell_input_fn` to get cell inputs.
        - Step 7: Call the wrapped `cell` with these cell inputs and
                  its previous state.
        - Step 8: (optional) Maybe copy output and cell state from history

        Args:
          inputs: (Possibly nested tuple of) Tensor,
                  the input at this time step.
          state: An instance of `TimeAttentionWrapperState`
                 containing tensors from the previous time step.

        Returns:
          A tuple `(attention_or_cell_output, next_state)`, where:

          - `attention_or_cell_output` depending on `output_attention`.
          - `next_state` is an instance of `TimeAttentionWrapperState`
             containing the state calculated at this time step.

        Raises:
          TypeError: If `state` is not an instance of
          `TimeAttentionWrapperState`.
        """
        if not isinstance(state, TimeAttentionWrapperState):
            raise TypeError("Expected state to be instance of "
                            "TimeAttentionWrapperState. "
                            "Received type {} instead.".format(type(state)))

        # Step 1: Calculate attention based on
        #         the previous output and current input
        cell_state = state.cell_state

        rnn_inputs, attn_inputs = self._rnn_and_attn_inputs_fn(inputs,
                                                               cell_state)

        cell_batch_size = (
                attn_inputs.shape[0].value or
                tf.shape(attn_inputs)[0])
        error_message = (
                "When applying AttentionWrapper %s: " % self.name +
                "Non-matching batch sizes between the memory "
                "(encoder output) and the query (decoder output).  "
                "Are you using "
                "the BeamSearchDecoder?  "
                "You may need to tile your memory input via "
                "the tf.contrib.seq2seq.tile_batch function with argument "
                "multiple=beam_width.")
        with tf.control_dependencies(
                self._batch_size_checks(cell_batch_size, error_message)):
            attn_inputs = tf.identity(
                    attn_inputs, name="checked_attn_inputs")

        if self._is_multi:
            previous_attention_state = state.attention_state
            previous_alignment_history = state.alignment_history
        else:
            previous_attention_state = [state.attention_state]
            previous_alignment_history = [state.alignment_history]

        all_alignments = []
        all_attentions = []
        all_attention_states = []
        maybe_all_histories = []

        prev_time_masks = self._read_from_tensor_array(state.all_time_masks,
                                                       state.time)
        prev_time_mask = prev_time_masks[:, -1, :]

        for i, attention_mechanism in enumerate(self._attention_mechanisms):
            # Steps 2 - 5 are performed inside `_compute_time_attention`
            (attention, alignments,
             next_attention_state) = _compute_time_attention(
                    attention_mechanism, attn_inputs,
                    previous_attention_state[i],
                    # time is added to calculate time attention
                    state.time, self._timed_ntms[i],
                    # provide boolean masks, to ignore some time steps
                    prev_time_mask, self._ignore_mask[i],
                    self._attention_layers[i]
                    if self._attention_layers else None)

            alignment_history = previous_alignment_history[i].write(
                    state.time, alignments) if self._alignment_history else ()

            all_attention_states.append(next_attention_state)
            all_alignments.append(alignments)
            all_attentions.append(attention)
            maybe_all_histories.append(alignment_history)

        attention = tf.concat(all_attentions, 1)

        # Step 6: Mix the `inputs` and `attention` output via
        #         `cell_input_fn` to get cell inputs.
        cell_inputs = self._cell_input_fn(rnn_inputs, attention)

        # Step 7: Call the wrapped `cell` with these cell inputs and
        #         its previous state.
        cell_output, next_cell_state = self._cell(cell_inputs, cell_state)

        prev_all_cell_states = state.all_cell_states

        time_mask = tf.concat([prev_time_mask[:, :state.time],
                               tf.ones_like(prev_time_mask[:, :1]),
                               prev_time_mask[:, state.time + 1:]], 1)

        if self._index_of_attn_to_copy is not None:
            # Step 8: Maybe copy output and cell state from history

            # get relevant previous outputs from history
            attn_to_copy = all_attentions[self._index_of_attn_to_copy]
            # copy them to current output
            cell_output_with_attn = cell_output + attn_to_copy

            memory_probs = self._get_memory_probs(all_alignments, state.time)

            # check that we do not pay attention to `tensor_not_to_copy`
            bin_likelihood_not_to_copy, _ = self._likelihood_fn(
                    cell_output_with_attn, self._tensor_not_to_copy)
            # recalculate probs
            memory_probs *= 1 - bin_likelihood_not_to_copy

            history_alignments = self._history_alignments(memory_probs)

            # get previous output from the history
            prev_output = self._prev_output(cell_output_with_attn,
                                            history_alignments,
                                            state.time)

            # check that current output is close to
            # the one in the history to which we pay attention to
            bin_likelihood_to_copy, _ = self._likelihood_fn(
                    cell_output_with_attn, prev_output)
            # recalculate probs
            memory_probs *= bin_likelihood_to_copy

            history_alignments = self._history_alignments(memory_probs)
            current_time_prob = history_alignments[:, -1:]

            # create additional likelihoods to maximize
            attn_likelihood = self._additional_likelihood(
                    attn_to_copy,
                    prev_output,
                    current_time_prob
            )
            state_likelihood = self._additional_likelihood(
                    cell_output + tf.stop_gradient(attn_to_copy),
                    prev_output,
                    current_time_prob
            )

            # recalculate time_mask
            time_mask = self._apply_alignments_to_history(
                    tf.cast(history_alignments, time_mask.dtype),
                    prev_time_masks[:, :-1, :],
                    time_mask
            )

            # recalculate new next_cell_state based on history_alignments
            next_cell_state = self._new_next_cell_state(
                    prev_all_cell_states,
                    next_cell_state,
                    cell_output_with_attn,
                    history_alignments,
                    state.time
            )

            all_cell_states = self._all_cell_states(
                    prev_all_cell_states,
                    next_cell_state,
                    state.time
            )

            if self._output_attention:
                # concatenate cell outputs, attention, additional likelihoods
                # and copy_attn_debug
                output = tf.concat([cell_output_with_attn,
                                    cell_output,
                                    attention,
                                    # additional likelihoods
                                    attn_likelihood, state_likelihood,
                                    # copy_attn_debug
                                    bin_likelihood_not_to_copy,
                                    bin_likelihood_to_copy,
                                    current_time_prob], 1)
            else:
                output = cell_output_with_attn

        else:
            # do not waste resources on storing history
            all_cell_states = prev_all_cell_states

            if self._output_attention:
                output = tf.concat([cell_output, attention], 1)
            else:
                output = cell_output

        all_time_masks = state.all_time_masks.write(state.time + 1, time_mask)

        next_state = TimeAttentionWrapperState(
                time=state.time + 1,
                cell_state=next_cell_state,
                attention=attention,
                attention_state=self._item_or_tuple(all_attention_states),
                alignments=self._item_or_tuple(all_alignments),
                alignment_history=self._item_or_tuple(maybe_all_histories),
                all_time_masks=all_time_masks,
                all_cell_states=all_cell_states
        )
        return output, next_state

    # helper for TensorArray
    @staticmethod
    def _read_from_tensor_array(tensor_array, time):
        """TensorArray time reader"""
        return tf.transpose(tensor_array.gather(tf.range(0, time + 1)),
                            [1, 0, 2])

    # helper methods for copy mechanism
    def _get_memory_probs(self, all_alignments, time):
        """Helper method to get memory_probs from all_alignments"""

        memory_probs = tf.stop_gradient(all_alignments[
                self._index_of_attn_to_copy][:, :time])

        # binarize memory_probs only if max value is larger than margin=0.1
        memory_probs_max = tf.reduce_max(memory_probs, axis=1, keepdims=True)
        memory_probs_max = tf.where(memory_probs_max > 0.1,
                                    memory_probs_max, -memory_probs_max)

        return tf.where(tf.equal(memory_probs, memory_probs_max),
                        tf.ones_like(memory_probs),
                        tf.zeros_like(memory_probs))

    @staticmethod
    def _history_alignments(memory_probs):
        """Helper method to apply binary mask to memory_probs"""

        current_time_prob = 1 - tf.reduce_sum(memory_probs, 1, keepdims=True)
        return tf.concat([memory_probs, current_time_prob], 1)

    @staticmethod
    def _apply_alignments_to_history(alignments, history_states, state):
        """Helper method to apply attention probabilities to rnn history

        copied from tf's `_compute_attention(...)`"""

        expanded_alignments = tf.stop_gradient(tf.expand_dims(alignments, 1))

        history_states = tf.concat([history_states,
                                    tf.expand_dims(state, 1)], 1)

        # Context is the inner product of alignments and values along the
        # memory time dimension.
        # expanded_alignments shape is
        #   [batch_size, 1, memory_time]
        # history_states shape is
        #   [batch_size, memory_time, memory_size]
        # the batched matmul is over memory_time, so the output shape is
        #   [batch_size, 1, memory_size].
        # we then squeeze out the singleton dim.

        return tf.squeeze(tf.matmul(expanded_alignments, history_states), [1])

    def _prev_output(self, state, alignments, time):
        """Helper method to get previous output from memory"""

        # get all previous outputs from appropriate
        # attention mechanism's memory limited by current time
        prev_outputs = tf.stop_gradient(self._attention_mechanisms[
                self._index_of_attn_to_copy].values[:, :time, :])

        # multiply by alignments to get one vector from one time step
        return self._apply_alignments_to_history(alignments,
                                                 prev_outputs,
                                                 state)

    def _additional_likelihood(self, output, prev_output, current_time_prob):
        """Helper method to create additional likelihood to maximize"""

        _, likelihood = self._likelihood_fn(
                output, tf.stop_gradient(prev_output))
        return tf.where(current_time_prob < 0.5,
                        likelihood, tf.ones_like(likelihood))

    def _new_hidden_state(self, prev_all_cell_states,
                          new_state, alignments, time):
        """Helper method to look into rnn history"""

        # reshape to (batch, time, memory_time) and
        # do not include current time because
        # we do not want to pay attention to it,
        # but we need to read it instead of
        # adding conditional flow if time == 0
        prev_cell_states = self._read_from_tensor_array(prev_all_cell_states,
                                                        time)[:, :-1, :]

        return self._apply_alignments_to_history(alignments,
                                                 prev_cell_states,
                                                 new_state)

    def _new_next_cell_state(self, prev_all_cell_states,
                             next_cell_state, new_cell_output,
                             alignments, time):
        """Helper method to recalculate new next_cell_state"""

        if isinstance(next_cell_state, tf.contrib.rnn.LSTMStateTuple):
            next_cell_state_c = self._new_hidden_state(
                prev_all_cell_states.c,
                next_cell_state.c,
                alignments,
                time
            )
            next_cell_state_h = self._new_hidden_state(
                prev_all_cell_states.h,
                new_cell_output,
                alignments,
                time
            )
            return tf.contrib.rnn.LSTMStateTuple(next_cell_state_c,
                                                 next_cell_state_h)
        else:
            return self._new_hidden_state(prev_all_cell_states,
                                          alignments, new_cell_output, time)

    @staticmethod
    def _all_cell_states(prev_all_cell_states, next_cell_state, time):
        """Helper method to recalculate all_cell_states tensor array"""

        if isinstance(next_cell_state, tf.contrib.rnn.LSTMStateTuple):
            return tf.contrib.rnn.LSTMStateTuple(
                    prev_all_cell_states.c.write(time + 1, next_cell_state.c),
                    prev_all_cell_states.h.write(time + 1, next_cell_state.h)
            )
        else:
            return prev_all_cell_states.write(time + 1, next_cell_state)


class ChronoBiasLayerNormBasicLSTMCell(tf.contrib.rnn.LayerNormBasicLSTMCell):
    """Custom LayerNormBasicLSTMCell that allows chrono initialization
        of gate biases.

        See super class for description."""

    def __init__(self,
                 num_units,
                 forget_bias=1.0,
                 input_bias=0.0,
                 activation=tf.tanh,
                 layer_norm=True,
                 norm_gain=1.0,
                 norm_shift=0.0,
                 dropout_keep_prob=1.0,
                 dropout_prob_seed=None,
                 out_layer_size=None,
                 reuse=None):
        """Initializes the basic LSTM cell

        Additional args:
            input_bias: float, The bias added to input gates.
            out_layer_size: (optional) integer, The number of units in
                the optional additional output layer.
        """
        super(ChronoBiasLayerNormBasicLSTMCell, self).__init__(
                num_units,
                forget_bias=forget_bias,
                activation=activation,
                layer_norm=layer_norm,
                norm_gain=norm_gain,
                norm_shift=norm_shift,
                dropout_keep_prob=dropout_keep_prob,
                dropout_prob_seed=dropout_prob_seed,
                reuse=reuse
        )
        self._input_bias = input_bias
        self._out_layer_size = out_layer_size

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

    @property
    def state_size(self):
        return tf.contrib.rnn.LSTMStateTuple(self._num_units,
                                             self.output_size)

    @staticmethod
    def _dense_layer(args, layer_size):
        """Optional out projection layer"""
        proj_size = args.get_shape()[-1]
        dtype = args.dtype
        weights = tf.get_variable("kernel",
                                  [proj_size, layer_size],
                                  dtype=dtype)
        bias = tf.get_variable("bias",
                               [layer_size],
                               dtype=dtype)
        out = tf.nn.bias_add(tf.matmul(args, weights), bias)
        return out

    def call(self, inputs, state):
        """LSTM cell with layer normalization and recurrent dropout."""
        c, h = state
        args = tf.concat([inputs, h], 1)
        concat = self._linear(args)
        dtype = args.dtype

        i, j, f, o = tf.split(value=concat, num_or_size_splits=4, axis=1)
        if self._layer_norm:
            i = self._norm(i, "input", dtype=dtype)
            j = self._norm(j, "transform", dtype=dtype)
            f = self._norm(f, "forget", dtype=dtype)
            o = self._norm(o, "output", dtype=dtype)

        g = self._activation(j)
        if (not isinstance(self._keep_prob, float)) or self._keep_prob < 1:
            g = tf.nn.dropout(g, self._keep_prob, seed=self._seed)

        new_c = (c * tf.sigmoid(f + self._forget_bias) +
                 g * tf.sigmoid(i + self._input_bias))  # added input_bias

        # do not do layer normalization on the new c,
        # because there are no trainable weights
        # if self._layer_norm:
        #     new_c = self._norm(new_c, "state", dtype=dtype)

        new_h = self._activation(new_c) * tf.sigmoid(o)

        # added dropout to the hidden state h
        if (not isinstance(self._keep_prob, float)) or self._keep_prob < 1:
            new_h = tf.nn.dropout(new_h, self._keep_prob, seed=self._seed)

        # add postprocessing of the output
        if self._out_layer_size is not None:
            with tf.variable_scope('out_layer'):
                new_h = self._dense_layer(new_h, self._out_layer_size)

        new_state = tf.contrib.rnn.LSTMStateTuple(new_c, new_h)
        return new_h, new_state
