# Copyright 2018 The TensorFlow Authors All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

from collections import namedtuple

import tensorflow as tf
import summary_utils as summ

Loss = namedtuple("Loss", "name loss vars")
Loss.__new__.__defaults__ = (tf.GraphKeys.TRAINABLE_VARIABLES,)


def iwae(model, observation, num_timesteps, num_samples=1,
         summarize=False):
  """Compute the IWAE evidence lower bound.

  Args:
    model: A callable that computes one timestep of the model.
    observation: A shape [batch_size*num_samples, state_size] Tensor
      containing z_n, the observation for each sequence in the batch.
    num_timesteps: The number of timesteps in each sequence, an integer.
    num_samples: The number of samples to use to compute the IWAE bound.
  Returns:
    log_p_hat: The IWAE estimator of the lower bound on the log marginal.
    loss: A tensor that you can perform gradient descent on to optimize the
      bound.
    maintain_ema_op: A no-op included for compatibility with FIVO.
    states: The sequence of states sampled.
  """
  # Initialization
  num_instances = tf.shape(observation)[0]
  batch_size = tf.cast(num_instances / num_samples, tf.int32)
  states = [model.zero_state(num_instances)]
  log_weights = []
  log_weight_acc = tf.zeros([num_samples, batch_size], dtype=observation.dtype)

  for t in xrange(num_timesteps):
    # run the model for one timestep
    (zt, log_q_zt, log_p_zt, log_p_x_given_z, _) = model(
        states[-1], observation, t)
    # update accumulators
    states.append(zt)
    log_weight = log_p_zt + log_p_x_given_z - log_q_zt
    log_weight_acc += tf.reshape(log_weight, [num_samples, batch_size])
    if summarize:
      weight_dist = tf.contrib.distributions.Categorical(
          logits=tf.transpose(log_weight_acc, perm=[1, 0]),
          allow_nan_stats=False)
      weight_entropy = weight_dist.entropy()
      weight_entropy = tf.reduce_mean(weight_entropy)
      tf.summary.scalar("weight_entropy/%d" % t, weight_entropy)
    log_weights.append(log_weight_acc)
  # Compute the lower bound on the log evidence.
  log_p_hat = (tf.reduce_logsumexp(log_weight_acc, axis=0) -
               tf.log(tf.cast(num_samples, observation.dtype))) / num_timesteps
  loss = -tf.reduce_mean(log_p_hat)
  losses = [Loss("log_p_hat", loss)]

  # we clip off the initial state before returning.
  # there are no emas for iwae, so we return a noop for that
  return log_p_hat, losses, tf.no_op(), states[1:], log_weights


def multinomial_resampling(log_weights, states, n, b):
  """Resample states with multinomial resampling.

  Args:
    log_weights: A (n x b) Tensor representing a batch of b logits for n-ary
      Categorical distribution.
    states: A list of (b*n x d) Tensors that will be resample in from the groups
     of every n-th row.

  Returns:
    resampled_states: A list of (b*n x d) Tensors resampled via stratified sampling.
    log_probs: A (n x b) Tensor of the log probabilities of the ancestry decisions.
    resampling_parameters: The Tensor of parameters of the resampling distribution.
    ancestors: An (n x b) Tensor of integral indices representing the ancestry decisions.
    resampling_dist: The distribution object for resampling.
  """
  log_weights = tf.convert_to_tensor(log_weights)
  states = [tf.convert_to_tensor(state) for state in states]

  resampling_parameters = tf.transpose(log_weights, perm=[1,0])
  resampling_dist = tf.contrib.distributions.Categorical(logits=resampling_parameters)
  ancestors = tf.stop_gradient(
      resampling_dist.sample(sample_shape=n))
  log_probs = resampling_dist.log_prob(ancestors)

  offset = tf.expand_dims(tf.range(b), 0)
  ancestor_inds = tf.reshape(ancestors * b + offset, [-1])

  resampled_states = []
  for state in states:
    resampled_states.append(tf.gather(state, ancestor_inds))
  return resampled_states, log_probs, resampling_parameters, ancestors, resampling_dist

def stratified_resampling(log_weights, states, n, b):
  """Resample states with straitified resampling.

  Args:
    log_weights: A (n x b) Tensor representing a batch of b logits for n-ary
      Categorical distribution.
    states: A list of (b*n x d) Tensors that will be resample in from the groups
     of every n-th row.

  Returns:
    resampled_states: A list of (b*n x d) Tensors resampled via stratified sampling.
    log_probs: A (n x b) Tensor of the log probabilities of the ancestry decisions.
    resampling_parameters: The Tensor of parameters of the resampling distribution.
    ancestors: An (n x b) Tensor of integral indices representing the ancestry decisions.
    resampling_dist: The distribution object for resampling.
  """
  log_weights = tf.convert_to_tensor(log_weights)
  states = [tf.convert_to_tensor(state) for state in states]

  log_weights = tf.transpose(log_weights, perm=[1,0])

  probs = tf.nn.softmax(
            tf.tile(tf.expand_dims(log_weights, axis=1),
                    [1, n, 1])
  )

  cdfs = tf.concat([tf.zeros((b,n,1), dtype=probs.dtype), tf.cumsum(probs, axis=2)], 2)

  bins = tf.range(n, dtype=probs.dtype) / n
  bins = tf.tile(tf.reshape(bins, [1,-1,1]), [b,1,n+1])

  strat_cdfs = tf.minimum(tf.maximum((cdfs - bins) * n, 0.0), 1.0)
  resampling_parameters = strat_cdfs[:,:,1:] - strat_cdfs[:,:,:-1]

  resampling_dist = tf.contrib.distributions.Categorical(
      probs = resampling_parameters,
      allow_nan_stats=False)

  ancestors = tf.stop_gradient(
      resampling_dist.sample())
  log_probs = resampling_dist.log_prob(ancestors)

  ancestors = tf.transpose(ancestors, perm=[1,0])
  log_probs = tf.transpose(log_probs, perm=[1,0])

  offset = tf.expand_dims(tf.range(b), 0)
  ancestor_inds = tf.reshape(ancestors * b + offset, [-1])

  resampled_states = []
  for state in states:
    resampled_states.append(tf.gather(state, ancestor_inds))

  return resampled_states, log_probs, resampling_parameters, ancestors, resampling_dist

def systematic_resampling(log_weights, states, n, b):
  """Resample states with systematic resampling.

  Args:
    log_weights: A (n x b) Tensor representing a batch of b logits for n-ary
      Categorical distribution.
    states: A list of (b*n x d) Tensors that will be resample in from the groups
     of every n-th row.

  Returns:
    resampled_states: A list of (b*n x d) Tensors resampled via stratified sampling.
    log_probs: A (n x b) Tensor of the log probabilities of the ancestry decisions.
    resampling_parameters: The Tensor of parameters of the resampling distribution.
    ancestors: An (n x b) Tensor of integral indices representing the ancestry decisions.
    resampling_dist: The distribution object for resampling.
  """

  log_weights = tf.convert_to_tensor(log_weights)
  states = [tf.convert_to_tensor(state) for state in states]

  log_weights = tf.transpose(log_weights, perm=[1,0])

  probs = tf.nn.softmax(
            tf.tile(tf.expand_dims(log_weights, axis=1),
                    [1, n, 1])
  )

  cdfs = tf.concat([tf.zeros((b,n,1), dtype=probs.dtype), tf.cumsum(probs, axis=2)], 2)

  bins = tf.range(n, dtype=probs.dtype) / n
  bins = tf.tile(tf.reshape(bins, [1,-1,1]), [b,1,n+1])

  strat_cdfs = tf.minimum(tf.maximum((cdfs - bins) * n, 0.0), 1.0)
  resampling_parameters = strat_cdfs[:,:,1:] - strat_cdfs[:,:,:-1]

  resampling_dist = tf.contrib.distributions.Categorical(
      probs=resampling_parameters,
      allow_nan_stats=True)

  U = tf.random_uniform((b, 1, 1), dtype=probs.dtype)

  ancestors = tf.stop_gradient(tf.reduce_sum(tf.to_float(U > strat_cdfs[:,:,1:]), axis=-1))
  log_probs = resampling_dist.log_prob(ancestors)

  ancestors = tf.transpose(ancestors, perm=[1,0])
  log_probs = tf.transpose(log_probs, perm=[1,0])

  offset = tf.expand_dims(tf.range(b, dtype=probs.dtype), 0)
  ancestor_inds = tf.reshape(ancestors * b + offset, [-1])

  resampled_states = []
  for state in states:
    resampled_states.append(tf.gather(state, ancestor_inds))

  return resampled_states, log_probs, resampling_parameters, ancestors, resampling_dist


def log_blend(inputs, weights):
  """Blends state in the log space.

  Args:
    inputs: A set of scalar states, one for each particle in each particle filter.
      Should be [num_samples, batch_size].
    weights: A set of weights used to blend the state. Each set of weights
      should be of dimension [num_samples] (one weight for each previous particle).
      There should be one set of weights for each new particle in each particle filter.
      Thus the shape should be [num_samples, batch_size, num_samples] where
      the first axis indexes new particle and the last axis indexes old particles.
  Returns:
    blended: The blended states, a tensor of shape [num_samples, batch_size].
  """
  raw_max = tf.reduce_max(inputs, axis=0, keepdims=True)
  my_max = tf.stop_gradient(
      tf.where(tf.is_finite(raw_max), raw_max, tf.zeros_like(raw_max))
  )
  # Don't ask.
  blended = tf.log(tf.einsum("ijk,kj->ij", weights, tf.exp(inputs - raw_max))) + my_max
  return blended


def relaxed_resampling(log_weights, states, num_samples, batch_size,
                       log_r_x=None, blend_type="log", temperature=0.5,
                       straight_through=False):
  """Resample states with relaxed resampling.

  Args:
    log_weights: A (n x b) Tensor representing a batch of b logits for n-ary
      Categorical distribution.
    states: A list of (b*n x d) Tensors that will be resample in from the groups
     of every n-th row.

  Returns:
    resampled_states: A list of (b*n x d) Tensors resampled via stratified sampling.
    log_probs: A (n x b) Tensor of the log probabilities of the ancestry decisions.
    resampling_parameters: The Tensor of parameters of the resampling distribution.
    ancestors: An (n x b x n) Tensor of relaxed one hot representations of the ancestry decisions.
    resampling_dist: The distribution object for resampling.
  """
  assert blend_type in ["log", "linear"], "Blend type must be 'log' or 'linear'."
  log_weights = tf.convert_to_tensor(log_weights)
  states = [tf.convert_to_tensor(state) for state in states]
  state_dim = states[0].get_shape().as_list()[-1]
  # weights are num_samples by batch_size, so we transpose to get a
  # set of batch_size distributions over [0,num_samples).
  resampling_parameters = tf.transpose(log_weights, perm=[1, 0])
  resampling_dist = tf.contrib.distributions.RelaxedOneHotCategorical(
      temperature,
      logits=resampling_parameters)

  # sample num_samples samples from the distribution, resulting in a
  # [num_samples, batch_size, num_samples] Tensor that represents a set of
  # [num_samples, batch_size] blending weights. The dimensions represent
  # [sample index, batch index, blending weight index]
  ancestors = resampling_dist.sample(sample_shape=num_samples)
  if straight_through:
    # Forward pass discrete choices, backwards pass soft choices
    hard_ancestor_indices = tf.argmax(ancestors, axis=-1)
    hard_ancestors = tf.one_hot(hard_ancestor_indices, num_samples,
                                dtype=ancestors.dtype)
    ancestors = tf.stop_gradient(hard_ancestors - ancestors) + ancestors
  log_probs = resampling_dist.log_prob(ancestors)
  if log_r_x is not None and blend_type == "log":
    log_r_x = tf.reshape(log_r_x, [num_samples, batch_size])
    log_r_x = log_blend(log_r_x, ancestors)
    log_r_x = tf.reshape(log_r_x, [num_samples*batch_size])
  elif log_r_x is not None and blend_type == "linear":
    # If blend type is linear just add log_r to the states that will be blended
    # linearly.
    states.append(log_r_x)

  # transpose the 'indices' to be [batch_index, blending weight index, sample index]
  ancestor_inds = tf.transpose(ancestors, perm=[1, 2, 0])
  resampled_states = []
  for state in states:
    # state is currently [num_samples * batch_size, state_dim] so we reshape
    # to [num_samples, batch_size, state_dim] and then transpose to
    # [batch_size, state_size, num_samples]
    state = tf.transpose(tf.reshape(state, [num_samples, batch_size, -1]), perm=[1, 2, 0])
    # state is now (batch_size, state_size, num_samples)
    # and ancestor is (batch index, blending weight index, sample index)
    # multiplying these gives a matrix of size [batch_size, state_size, num_samples]
    next_state = tf.matmul(state, ancestor_inds)
    # transpose the state to be [num_samples, batch_size, state_size]
    # and then reshape it to match the state format.
    next_state = tf.reshape(tf.transpose(next_state, perm=[2,0,1]), [num_samples*batch_size, state_dim])
    resampled_states.append(next_state)

  new_dist = tf.contrib.distributions.Categorical(
      logits=resampling_parameters)

  if log_r_x is not None and blend_type == "linear":
    # If blend type is linear pop off log_r that we added to the states.
    log_r_x = tf.squeeze(resampled_states[-1])
    resampled_states = resampled_states[:-1]
  return resampled_states, log_probs, log_r_x, resampling_parameters, ancestors, new_dist


def fivo(model,
         observation,
         num_timesteps,
         resampling_schedule,
         num_samples=1,
         use_resampling_grads=True,
         resampling_type="multinomial",
         resampling_temperature=0.5,
         aux=True,
         summarize=False):
  """Compute the FIVO evidence lower bound.

  Args:
    model: A callable that computes one timestep of the model.
    observation: A shape [batch_size*num_samples, state_size] Tensor
      containing z_n, the observation for each sequence in the batch.
    num_timesteps: The number of timesteps in each sequence, an integer.
    resampling_schedule: A list of booleans of length num_timesteps, contains
      True if a resampling should occur on a specific timestep.
    num_samples: The number of samples to use to compute the IWAE bound.
    use_resampling_grads: Whether or not to include the resampling gradients
      in loss.
    resampling type: The type of resampling, one of "multinomial", "stratified",
      "relaxed-logblend", "relaxed-linearblend", "relaxed-stateblend", or
      "systematic".
    resampling_temperature: A positive temperature only used for relaxed
      resampling.
    aux: If true, compute the FIVO-AUX bound.
  Returns:
    log_p_hat: The IWAE estimator of the lower bound on the log marginal.
    loss: A tensor that you can perform gradient descent on to optimize the
      bound.
    maintain_ema_op: An op to update the baseline ema used for the resampling
      gradients.
    states: The sequence of states sampled.
  """
  # Initialization
  num_instances = tf.cast(tf.shape(observation)[0], tf.int32)
  batch_size = tf.cast(num_instances / num_samples, tf.int32)
  states = [model.zero_state(num_instances)]
  prev_state = states[0]
  log_weight_acc = tf.zeros(shape=[num_samples, batch_size], dtype=observation.dtype)
  prev_log_r_zt = tf.zeros([num_instances], dtype=observation.dtype)
  log_weights = []
  log_weights_all = []
  log_p_hats = []
  resampling_log_probs = []
  for t in xrange(num_timesteps):
    # run the model for one timestep
    (zt, log_q_zt, log_p_zt, log_p_x_given_z, log_r_zt) = model(
        prev_state, observation, t)
    # update accumulators
    states.append(zt)
    log_weight = log_p_zt + log_p_x_given_z - log_q_zt
    if aux:
      if t == num_timesteps - 1:
        log_weight -= prev_log_r_zt
      else:
        log_weight += log_r_zt - prev_log_r_zt
      prev_log_r_zt = log_r_zt
    log_weight_acc += tf.reshape(log_weight, [num_samples, batch_size])
    log_weights_all.append(log_weight_acc)
    if resampling_schedule[t]:

      # These objects will be resampled
      to_resample = [states[-1]]
      if aux and "relaxed" not in resampling_type:
        to_resample.append(prev_log_r_zt)

      # do the resampling
      if resampling_type == "multinomial":
        (resampled,
         resampling_log_prob,
         _, _, _) = multinomial_resampling(log_weight_acc,
                                           to_resample,
                                           num_samples,
                                           batch_size)
      elif resampling_type == "stratified":
        (resampled,
         resampling_log_prob,
         _, _, _) = stratified_resampling(log_weight_acc,
                                          to_resample,
                                          num_samples,
                                          batch_size)
      elif resampling_type == "systematic":
        (resampled,
         resampling_log_prob,
         _, _, _) = systematic_resampling(log_weight_acc,
                                          to_resample,
                                          num_samples,
                                          batch_size)
      elif "relaxed" in resampling_type:
        if aux:
          if resampling_type == "relaxed-logblend":
            (resampled,
             resampling_log_prob,
             prev_log_r_zt,
             _, _, _) = relaxed_resampling(log_weight_acc,
                                           to_resample,
                                           num_samples,
                                           batch_size,
                                           temperature=resampling_temperature,
                                           log_r_x=prev_log_r_zt,
                                           blend_type="log")
          elif resampling_type == "relaxed-linearblend":
            (resampled,
             resampling_log_prob,
             prev_log_r_zt,
             _, _, _) = relaxed_resampling(log_weight_acc,
                                           to_resample,
                                           num_samples,
                                           batch_size,
                                           temperature=resampling_temperature,
                                           log_r_x=prev_log_r_zt,
                                           blend_type="linear")
          elif resampling_type == "relaxed-stateblend":
            (resampled,
             resampling_log_prob,
             _, _, _, _) = relaxed_resampling(log_weight_acc,
                                              to_resample,
                                              num_samples,
                                              batch_size,
                                              temperature=resampling_temperature)
            # Calculate prev_log_r_zt from the post-resampling state
            prev_r_zt = model.r.r_xn(resampled[0], t)
            prev_log_r_zt = tf.reduce_sum(
                prev_r_zt.log_prob(observation), axis=[1])
          elif resampling_type == "relaxed-stateblend-st":
            (resampled,
             resampling_log_prob,
             _, _, _, _) = relaxed_resampling(log_weight_acc,
                                              to_resample,
                                              num_samples,
                                              batch_size,
                                              temperature=resampling_temperature,
                                              straight_through=True)
            # Calculate prev_log_r_zt from the post-resampling state
            prev_r_zt = model.r.r_xn(resampled[0], t)
            prev_log_r_zt = tf.reduce_sum(
                prev_r_zt.log_prob(observation), axis=[1])
        else:
          (resampled,
           resampling_log_prob,
           _, _, _, _) = relaxed_resampling(log_weight_acc,
                                            to_resample,
                                            num_samples,
                                            batch_size,
                                            temperature=resampling_temperature)
      #if summarize:
      #  resampling_entropy = resampling_dist.entropy()
      #  resampling_entropy = tf.reduce_mean(resampling_entropy)
      #  tf.summary.scalar("weight_entropy/%d" % t, resampling_entropy)

      resampling_log_probs.append(tf.reduce_sum(resampling_log_prob, axis=0))
      prev_state = resampled[0]
      if aux and "relaxed" not in resampling_type:
        # Squeeze out the extra dim potentially added by resampling.
        # prev_log_r_zt should always be [num_instances]
        prev_log_r_zt = tf.squeeze(resampled[1])
      # Update the log p hat estimate, taking a log sum exp over the sample
      # dimension. The appended tensor is [batch_size].
      log_p_hats.append(
          tf.reduce_logsumexp(log_weight_acc, axis=0) - tf.log(
              tf.cast(num_samples, dtype=observation.dtype)))
      # reset the weights
      log_weights.append(log_weight_acc)
      log_weight_acc = tf.zeros_like(log_weight_acc)
    else:
      prev_state = states[-1]
  # Compute the final weight update. If we just resampled this will be zero.
  final_update = (tf.reduce_logsumexp(log_weight_acc, axis=0) -
                  tf.log(tf.cast(num_samples, dtype=observation.dtype)))
  # If we ever resampled, then sum up the previous log p hat terms
  if len(log_p_hats) > 0:
    log_p_hat = tf.reduce_sum(log_p_hats, axis=0) + final_update
  else:  # otherwise, log_p_hat only comes from the final update
    log_p_hat = final_update

  if use_resampling_grads and any(resampling_schedule):
    # compute the rewards
    # cumsum([a, b, c]) => [a, a+b, a+b+c]
    # learning signal at timestep t is
    #   [sum from i=t+1 to T of log_p_hat_i for t=1:T]
    # so we will compute (sum from i=1 to T of log_p_hat_i)
    # and at timestep t will subtract off (sum from i=1 to t of log_p_hat_i)
    # rewards is a [num_resampling_events, batch_size] Tensor
    rewards = tf.stop_gradient(
        tf.expand_dims(log_p_hat, 0) - tf.cumsum(log_p_hats, axis=0))
    batch_avg_rewards = tf.reduce_mean(rewards, axis=1)
    # compute ema baseline.
    # centered_rewards is [num_resampling_events, batch_size]
    baseline_ema = tf.train.ExponentialMovingAverage(decay=0.94)
    maintain_baseline_op = baseline_ema.apply([batch_avg_rewards])
    baseline = tf.expand_dims(baseline_ema.average(batch_avg_rewards), 1)
    centered_rewards = rewards - baseline
    if summarize:
      summ.summarize_learning_signal(rewards, "rewards")
      summ.summarize_learning_signal(centered_rewards, "centered_rewards")
    # compute the loss tensor.
    resampling_grads = tf.reduce_sum(
        tf.stop_gradient(centered_rewards) * resampling_log_probs, axis=0)
    losses = [Loss("log_p_hat", -tf.reduce_mean(log_p_hat)/num_timesteps),
              Loss("resampling_grads", -tf.reduce_mean(resampling_grads)/num_timesteps)]
  else:
    losses = [Loss("log_p_hat", -tf.reduce_mean(log_p_hat)/num_timesteps)]
    maintain_baseline_op = tf.no_op()

  log_p_hat /= num_timesteps
  # we clip off the initial state before returning.
  return log_p_hat, losses, maintain_baseline_op, states[1:], log_weights_all


def fivo_aux_td(
    model,
    observation,
    num_timesteps,
    resampling_schedule,
    num_samples=1,
    summarize=False):
  """Compute the FIVO_AUX evidence lower bound."""
  # Initialization
  num_instances = tf.cast(tf.shape(observation)[0], tf.int32)
  batch_size = tf.cast(num_instances / num_samples, tf.int32)
  states = [model.zero_state(num_instances)]
  prev_state = states[0]
  log_weight_acc = tf.zeros(shape=[num_samples, batch_size], dtype=observation.dtype)
  prev_log_r = tf.zeros([num_instances], dtype=observation.dtype)
  # must be pre-resampling
  log_rs = []
  # must be post-resampling
  r_tilde_params = [model.r_tilde.r_zt(states[0], observation, 0)]
  log_r_tildes = []
  log_p_xs = []
  # contains the weight at each timestep before resampling only on resampling timesteps
  log_weights = []
  # contains weight at each timestep before resampling
  log_weights_all = []
  log_p_hats = []
  for t in xrange(num_timesteps):
    # run the model for one timestep
    # zt is state, [num_instances, state_dim]
    # log_q_zt, log_p_x_given_z is [num_instances]
    # r_tilde_mu, r_tilde_sigma is [num_instances, state_dim]
    # p_ztplus1 is a normal distribution on [num_instances, state_dim]
    (zt, log_q_zt, log_p_zt, log_p_x_given_z,
     r_tilde_mu, r_tilde_sigma_sq, p_ztplus1) = model(prev_state, observation, t)

    # Compute the log weight without log r.
    log_weight = log_p_zt + log_p_x_given_z - log_q_zt

    # Compute log r.
    if t == num_timesteps - 1:
      log_r = tf.zeros_like(prev_log_r)
    else:
      p_mu = p_ztplus1.mean()
      p_sigma_sq = p_ztplus1.variance()
      log_r = (tf.log(r_tilde_sigma_sq) -
               tf.log(r_tilde_sigma_sq + p_sigma_sq) -
               tf.square(r_tilde_mu - p_mu)/(r_tilde_sigma_sq + p_sigma_sq))
      log_r = 0.5*tf.reduce_sum(log_r, axis=-1)

    #log_weight += tf.stop_gradient(log_r - prev_log_r)
    log_weight += log_r - prev_log_r
    log_weight_acc += tf.reshape(log_weight, [num_samples, batch_size])

    # Update accumulators
    states.append(zt)
    log_weights_all.append(log_weight_acc)
    log_p_xs.append(log_p_x_given_z)
    log_rs.append(log_r)

    # Compute log_r_tilde as [num_instances] Tensor.
    prev_r_tilde_mu, prev_r_tilde_sigma_sq = r_tilde_params[-1]
    prev_log_r_tilde = -0.5*tf.reduce_sum(
        tf.square(zt - prev_r_tilde_mu)/prev_r_tilde_sigma_sq, axis=-1)
        #tf.square(tf.stop_gradient(zt) - r_tilde_mu)/r_tilde_sigma_sq, axis=-1)
        #tf.square(zt - r_tilde_mu)/r_tilde_sigma_sq, axis=-1)
    log_r_tildes.append(prev_log_r_tilde)

    # optionally resample
    if resampling_schedule[t]:
      # These objects will be resampled
      if t < num_timesteps - 1:
        to_resample = [zt, log_r, r_tilde_mu, r_tilde_sigma_sq]
      else:
        to_resample = [zt, log_r]
      (resampled,
       _, _, _, _) = multinomial_resampling(log_weight_acc,
                                            to_resample,
                                            num_samples,
                                            batch_size)
      prev_state = resampled[0]
      # Squeeze out the extra dim potentially added by resampling.
      # prev_log_r_zt and log_r_tilde should always be [num_instances]
      prev_log_r = tf.squeeze(resampled[1])
      if t < num_timesteps -1:
        r_tilde_params.append((resampled[2], resampled[3]))
      # Update the log p hat estimate, taking a log sum exp over the sample
      # dimension. The appended tensor is [batch_size].
      log_p_hats.append(
          tf.reduce_logsumexp(log_weight_acc, axis=0) - tf.log(
              tf.cast(num_samples, dtype=observation.dtype)))
      # reset the weights
      log_weights.append(log_weight_acc)
      log_weight_acc = tf.zeros_like(log_weight_acc)
    else:
      prev_state = zt
      prev_log_r = log_r
      if t < num_timesteps - 1:
        r_tilde_params.append((r_tilde_mu, r_tilde_sigma_sq))

  # Compute the final weight update. If we just resampled this will be zero.
  final_update = (tf.reduce_logsumexp(log_weight_acc, axis=0) -
                  tf.log(tf.cast(num_samples, dtype=observation.dtype)))
  # If we ever resampled, then sum up the previous log p hat terms
  if len(log_p_hats) > 0:
    log_p_hat = tf.reduce_sum(log_p_hats, axis=0) + final_update
  else:  # otherwise, log_p_hat only comes from the final update
    log_p_hat = final_update

  # Compute the bellman loss.
  # Will remove the first timestep as it is not used.
  # log p(x_t|z_t) is in row t-1.
  log_p_x = tf.reshape(tf.stack(log_p_xs),
                       [num_timesteps, num_samples, batch_size])
  # log r_t is contained in row t-1.
  # last column is zeros (because at timestep T (num_timesteps) r is 1.
  log_r = tf.reshape(tf.stack(log_rs),
                     [num_timesteps, num_samples, batch_size])
  # [num_timesteps, num_instances]. log r_tilde_t is in row t-1.
  log_r_tilde = tf.reshape(tf.stack(log_r_tildes),
                           [num_timesteps, num_samples, batch_size])
  log_lambda = tf.reduce_mean(log_r_tilde - log_p_x - log_r, axis=1,
                              keepdims=True)
  bellman_sos = tf.reduce_mean(tf.square(
      log_r_tilde - tf.stop_gradient(log_lambda + log_p_x + log_r)), axis=[0, 1])
  bellman_loss = tf.reduce_mean(bellman_sos)/num_timesteps
  tf.summary.scalar("bellman_loss", bellman_loss)

  if len(tf.get_collection("LOG_P_HAT_VARS")) == 0:
    log_p_hat_collection = list(set(tf.trainable_variables()) -
                                set(tf.get_collection("R_TILDE_VARS")))
    for v in log_p_hat_collection:
      tf.add_to_collection("LOG_P_HAT_VARS", v)

  log_p_hat /= num_timesteps
  losses = [Loss("log_p_hat", -tf.reduce_mean(log_p_hat), "LOG_P_HAT_VARS"),
            Loss("bellman_loss", bellman_loss, "R_TILDE_VARS")]

  return log_p_hat, losses, tf.no_op(), states[1:], log_weights_all