# Copyright 2018 The TensorFlow Probability Authors.
#
# 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.
# ============================================================================
"""The differential evolution global optimization algorithm.

Differential evolution (DE) is a population based global optimization scheme.
It is applicable to problems with continuous parameter space and because it
does not require gradient computation, it is also applicable to
non-differentiable functions. For more details see:
https://en.wikipedia.org/wiki/Differential_evolution

DE starts with a population of candidate solutions (represented as vectors).
It generates new trial solutions by adding the weighted difference between
two population vectors to a third vector. This is the mutation operation. The
mutated vectors parameters are mixed with the parameters of a third
vector to yield the final candidate solution. This is the crossover step. If
the trial vector yields a lower cost value it replaces the target vector. This
process is repeated for each member of the population to complete one
generation.

There are a number of different schemes that fall under the DE umbrella.
The established notation for representing these schemes
is as `DE/x/y/z` where `x` specifies how the population member to be mutated is
selected (may be `rand` if it is chosen randomly or `best` if the best
member is chosen), `y` specifies the number of difference vectors used and `z`
denotes the crossover scheme employed. This may be `bin` if binary
recombination is used or `exp` if exponential crossover is used.

The most commonly employed scheme is `DE/rand/1/bin`. This is the one
implemented in this module.
"""

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

import collections
import numpy as np
import tensorflow as tf

from tensorflow_probability.python import distributions

_DifferentialEvolutionOptimizerResults = collections.namedtuple(
    '_DifferentialEvolutionOptimizerResults', [
        'converged',  # Scalar boolean tensor indicating whether the minimum
                      # was found within tolerance.
        'failed',  # Scalar boolean tensor indicating whether the search failed.
                   # This may happen if the objective values become NaN for
                   # the entire population.
        'position',  # A list of tensors containing the best argument value
                     # found during the search. If the search converged, then
                     # this value is the argmin of the objective function.
        'objective_value',  # A tensor containing the value of the objective
                            # function at the `position`. If the search
                            # converged, then this is the (local) minimum of
                            # the objective function.
        'final_population',  # The final state of the population.
        'final_objective_values',  # The objective function evaluated at the
                                   # final population.
        'initial_population',  # The starting population.
        'initial_objective_values',  # The objective function evaluated at the
                                     # initial population.
        'num_iterations'  # The number of generations the population was
                          # evolved.
    ])

class DifferentialEvolutionOptimizerResults(
    _DifferentialEvolutionOptimizerResults):
  """Results of a differential evolution optimization run.

  The object has the following attributes:
    converged: Scalar boolean `Tensor` indicating whether the minimum was
      found within the specified tolerances.
    failed: Scalar boolean tensor indicating whether the search failed.
      This may happen if the objective values become NaN for the entire
      population.
    position: A `Tensor` containing the best point found during the search.
      If the search converged, then this value is the argmin of the
      objective function within the specified tolerances.
    objective_value: A tensor containing the value of the objective
      function at the `position`. If the search
      converged, then this is the (local) minimum of
      the objective function.
    final_population: The final state of the population.
    final_objective_values: The objective function evaluated at the
      final population.
    initial_population: The starting population.
    initial_objective_values: The objective function evaluated at the
      initial population.
    num_iterations: The number of generations the population was evolved.
  """

# Class to keep track of the loop variables in the minimize method.
_MinimizeLoopVars = collections.namedtuple(
    '_MinimizeLoopVars',
    [
        'converged',
        'failed',
        'num_iterations',
        'population',
        'population_values'
    ])

def one_step(
    objective_function,
    population,
    population_values=None,
    differential_weight=0.5,
    crossover_prob=0.9,
    seed=None,
    name=None):
  """Performs one step of the differential evolution algorithm.

  Args:
    objective_function:  A Python callable that accepts a batch of possible
      solutions and returns the values of the objective function at those
      arguments as a rank 1 real `Tensor`. This specifies the function to be
      minimized. The input to this callable may be either a single `Tensor`
      or a Python `list` of `Tensor`s. The signature must match the format of
      the argument `population`. (i.e. objective_function(*population) must
      return the value of the function to be minimized).
    population:  `Tensor` or Python `list` of `Tensor`s representing the
      current population vectors. Each `Tensor` must be of the same real dtype.
      The first dimension indexes individual population members while the
      rest of the dimensions are consumed by the value function. For example,
      if the population is a single `Tensor` of shape [n, m1, m2], then `n` is
      the population size and the output of `objective_function` applied to the
      population is a `Tensor` of shape [n]. If the population is a python
      list of `Tensor`s then each `Tensor` in the list should have the first
      axis of a common size, say `n` and `objective_function(*population)`
      should return a `Tensor of shape [n]. The population must have at least
      4 members for the algorithm to work correctly.
    population_values: A `Tensor` of rank 1 and real dtype. The result of
      applying `objective_function` to the `population`. If not supplied it is
      computed using the `objective_function`.
      Default value: None.
    differential_weight: Real scalar `Tensor`. Must be positive and less than
      2.0. The parameter controlling the strength of mutation.
      Default value: 0.5
    crossover_prob: Real scalar `Tensor`. Must be between 0 and 1. The
      probability of recombination per site.
      Default value: 0.9
    seed: `int` or None. The random seed for this `Op`. If `None`, no seed is
      applied.
      Default value: None.
    name: (Optional) Python str. The name prefixed to the ops created by this
      function. If not supplied, the default name 'one_step' is
      used.
      Default value: None

  Returns:
    A sequence containing the following elements (in order):
    next_population: A `Tensor` or Python `list` of `Tensor`s of the same
      structure as the input population. The population at the next generation.
    next_population_values: A `Tensor` of same shape and dtype as input
      `population_values`. The function values for the `next_population`.
  """
  with tf.compat.v1.name_scope(
      name, 'one_step',
      [population, population_values, differential_weight, crossover_prob]):
    population, _ = _ensure_list(population)
    if population_values is None:
      population_values = objective_function(*population)
    population_size = tf.shape(input=population[0])[0]
    seed_stream = distributions.SeedStream(seed, salt='one_step')
    mixing_indices = _get_mixing_indices(population_size, seed=seed_stream())
    # Construct the mutated solution vectors. There is one for each member of
    # the population.
    mutants = _get_mutants(population,
                           population_size,
                           mixing_indices,
                           differential_weight)
    # Perform recombination between the parents and the mutants.
    candidates = _binary_crossover(population,
                                   population_size,
                                   mutants,
                                   crossover_prob,
                                   seed=seed_stream())
    candidate_values = objective_function(*candidates)
    if population_values is None:
      population_values = objective_function(*population)

    infinity = tf.zeros_like(population_values) + np.inf

    population_values = tf.where(
        tf.math.is_nan(population_values), x=infinity, y=population_values)

    to_replace = candidate_values < population_values
    next_population = [
        tf.where(to_replace, x=candidates_part, y=population_part)
        for candidates_part, population_part in zip(candidates, population)
    ]
    next_values = tf.where(to_replace, x=candidate_values, y=population_values)

  return next_population, next_values

def minimize(objective_function,
             initial_population=None,
             initial_position=None,
             population_size=50,
             population_stddev=1.,
             max_iterations=100,
             func_tolerance=0,
             position_tolerance=1e-8,
             differential_weight=0.5,
             crossover_prob=0.9,
             seed=None,
             name=None):
  """Applies the Differential evolution algorithm to minimize a function.

  Differential Evolution is an evolutionary optimization algorithm which works
  on a set of candidate solutions called the population. It iteratively
  improves the population by applying genetic operators of mutation and
  recombination. The objective function `f` supplies the fitness of each
  candidate. A candidate `s_1` is considered better than `s_2` if
  `f(s_1) < f(s_2)`.

  This method allows the user to either specify an initial population or a
  single candidate solution. If a single solution is specified, a population
  of the specified size is initialized by adding independent normal noise
  to the candidate solution.

  The implementation also supports a multi-part specification of the state. For
  example, consider the objective function:

  ```python
  # x is a tensor of shape [n, m] while y is of shape [n].
  def objective(x, y):
    return tf.math.reduce_sum(x ** 2, axis=-1) + y ** 2
  ```
  The state in this case is specified by two input tensors `x` and `y`. To
  apply the algorithm to this objective function, one would need to specify
  either an initial population as a list of two tensors of shapes
  `[population_size, k]` and `[population_size]`. The following code shows the
  complete example:

  ```python
    population_size = 40
    # With an initial population and a multi-part state.
    initial_population = (tf.random.normal([population_size]),
                          tf.random.normal([population_size]))
    def easom_fn(x, y):
      return -(tf.math.cos(x) * tf.math.cos(y) *
               tf.math.exp(-(x-np.pi)**2 - (y-np.pi)**2))

    optim_results = tfp.optimizers.differential_evolution_minimize(
        easom_fn,
        initial_population=initial_population,
        seed=43210)

    print (optim_results.converged)
    print (optim_results.position)  # Should be (close to) [pi, pi].
    print (optim_results.objective_value)    # Should be -1.

    # With a single starting point
    initial_position = (tf.constant(1.0), tf.constant(1.0))

    optim_results = tfp.optimizers.differential_evolution_minimize(
        easom_fn,
        initial_position=initial_position,
        population_size=40,
        population_stddev=2.0,
        seed=43210)
  ```

  Args:
    objective_function: A Python callable that accepts a batch of possible
      solutions and returns the values of the objective function at those
      arguments as a rank 1 real `Tensor`. This specifies the function to be
      minimized. The input to this callable may be either a single `Tensor`
      or a Python `list` of `Tensor`s. The signature must match the format of
      the argument `population`. (i.e. objective_function(*population) must
      return the value of the function to be minimized).
    initial_population: A real `Tensor` or Python list of `Tensor`s.
      If a list, each `Tensor` must be of rank at least 1 and with a common
      first dimension. The first dimension indexes into the candidate solutions
      while the rest of the dimensions (if any) index into an individual
      solution. The size of the population must be at least 4. This is a
      requirement of the DE algorithm.
    initial_position: A real `Tensor` of any shape. The seed solution used
      to initialize the population of solutions. If this parameter is specified
      then `initial_population` must not be specified.
    population_size: A positive scalar int32 `Tensor` greater than 4. The
      size of the population to evolve. This parameter is ignored if
      `initial_population` is specified.
      Default value: 50.
    population_stddev: A positive scalar real `Tensor` of the same dtype
      as `initial_position`. This parameter is ignored if `initial_population`
      is specified. Used to generate the population from the `initial_position`
      by adding random normal noise with zero mean and the specified standard
      deviation.
      Default value: 1.0
    max_iterations: Positive scalar int32 `Tensor`. The maximum number of
      generations to evolve the population for.
      Default value: 100
    func_tolerance: Scalar `Tensor` of the same dtype as the output of the
      `objective_function`. The algorithm stops if the absolute difference
      between the largest and the smallest objective function value in the
      population is below this number.
      Default value: 0
    position_tolerance: Scalar `Tensor` of the same real dtype as
      `initial_position` or `initial_population`. The algorithm terminates if
      the largest absolute difference between the coordinates of the population
      members is below this threshold.
      Default value: 1e-8
    differential_weight: Real scalar `Tensor`. Must be positive and less than
      2.0. The parameter controlling the strength of mutation in the algorithm.
      Default value: 0.5
    crossover_prob: Real scalar `Tensor`. Must be between 0 and 1. The
      probability of recombination per site.
      Default value: 0.9
    seed: `int` or None. The random seed for this `Op`. If `None`, no seed is
      applied.
      Default value: None.
    name: (Optional) Python str. The name prefixed to the ops created by this
      function. If not supplied, the default name
      'differential_evolution_minimize' is used.
      Default value: None

  Returns:
    optimizer_results: An object containing the following attributes:
      converged: Scalar boolean `Tensor` indicating whether the minimum was
        found within the specified tolerances.
      num_objective_evaluations: The total number of objective
        evaluations performed.
      position: A `Tensor` containing the best point found during the search.
        If the search converged, then this value is the argmin of the
        objective function within the specified tolerances.
      objective_value: A `Tensor` containing the value of the objective
        function at the `position`. If the search
        converged, then this is the (local) minimum of
        the objective function.
      final_population: The final state of the population.
      final_objective_values: The objective function evaluated at the
        final population.
      initial_population: The starting population.
      initial_objective_values: The objective function evaluated at the
        initial population.
      num_iterations: The number of iterations of the main algorithm body.

  Raises:
    ValueError: If neither the initial population, nor the initial position
      are specified or if both are specified.
  """

  if initial_population is None and initial_position is None:
    raise ValueError('Either the initial population or the initial position '
                     'must be specified.')
  if initial_population is not None and initial_position is not None:
    raise ValueError('Only one of initial population or initial position '
                     'should be specified')

  with tf.compat.v1.name_scope(
      name,
      default_name='minimize',
      values=[
          initial_population, initial_position, population_size,
          population_stddev, max_iterations, func_tolerance, position_tolerance,
          differential_weight, crossover_prob
      ]):
    (
        was_iterable,
        population,
        population_values,
        max_iterations,
        func_tolerance,
        position_tolerance,
        differential_weight,
        crossover_prob
    ) = _get_initial_args(objective_function,
                          initial_population,
                          initial_position,
                          population_size,
                          population_stddev,
                          max_iterations,
                          func_tolerance,
                          position_tolerance,
                          differential_weight,
                          crossover_prob,
                          seed)

    def evolve_body(loop_vars):
      """Performs one step of the evolution."""
      next_population, next_population_values = one_step(
          objective_function,
          loop_vars.population,
          population_values=loop_vars.population_values,
          differential_weight=differential_weight,
          crossover_prob=crossover_prob,
          seed=seed)
      converged = _check_convergence(next_population,
                                     next_population_values,
                                     func_tolerance,
                                     position_tolerance)

      failed = _check_failure(next_population_values)

      return [_MinimizeLoopVars(
          converged=converged,
          failed=failed,
          num_iterations=loop_vars.num_iterations+1,
          population=next_population,
          population_values=next_population_values)]

    def evolve_cond(loop_vars):
      should_stop = (
          loop_vars.failed |
          loop_vars.converged |
          (max_iterations is not None and
           loop_vars.num_iterations >= max_iterations))
      return ~should_stop

    initial_vars = _MinimizeLoopVars(
        converged=tf.convert_to_tensor(value=False),
        failed=tf.convert_to_tensor(value=False),
        num_iterations=tf.convert_to_tensor(value=0),
        population=population,
        population_values=population_values)
    final_state = tf.while_loop(
        cond=evolve_cond, body=evolve_body, loop_vars=(initial_vars,))[0]
    best_position, best_values = _find_best_in_population(
        final_state.population,
        final_state.population_values)
    # Ensure we return a similar structure to what the user supplied.
    final_population = final_state.population
    if not was_iterable:
      final_population = final_population[0]
      best_position = best_position[0]
    return DifferentialEvolutionOptimizerResults(
        converged=final_state.converged,
        failed=final_state.failed,
        position=best_position,
        objective_value=best_values,
        final_population=final_population,
        final_objective_values=final_state.population_values,
        initial_population=population,
        initial_objective_values=population_values,
        num_iterations=final_state.num_iterations)

def _get_initial_args(objective_function,
                      initial_population,
                      initial_position,
                      population_size,
                      population_stddev,
                      max_iterations,
                      func_tolerance,
                      position_tolerance,
                      differential_weight,
                      crossover_prob,
                      seed):
  """Processes initial args."""
  was_iterable = False
  if initial_position is not None:
    initial_position, was_iterable = _ensure_list(initial_position)

  if initial_population is not None:
    initial_population, was_iterable = _ensure_list(initial_population)

  population = _get_starting_population(initial_population,
                                        initial_position,
                                        population_size,
                                        population_stddev,
                                        seed=seed)

  differential_weight = tf.convert_to_tensor(
      value=differential_weight, dtype=population[0].dtype.base_dtype)

  crossover_prob = tf.convert_to_tensor(value=crossover_prob)
  population_values = objective_function(*population)
  if max_iterations is not None:
    max_iterations = tf.convert_to_tensor(value=max_iterations)
  func_tolerance = tf.convert_to_tensor(
      value=func_tolerance, dtype=population_values.dtype.base_dtype)
  position_tolerance = tf.convert_to_tensor(
      value=position_tolerance, dtype=population[0].dtype.base_dtype)
  return (was_iterable,
          population,
          population_values,
          max_iterations,
          func_tolerance,
          position_tolerance,
          differential_weight,
          crossover_prob)

def _check_failure(population_values):
  """Checks if all the population values are NaN/infinite."""
  return tf.math.reduce_all(input_tensor=tf.math.is_inf(population_values))

def _find_best_in_population(population, values):
  """Finds the population member with the lowest value."""
  best_value = tf.math.reduce_min(input_tensor=values)
  best_index = tf.where(tf.math.equal(values, best_value))[0, 0]

  return ([population_part[best_index] for population_part in population],
          best_value)

def _check_convergence(population,
                       population_values,
                       func_tolerance,
                       position_tolerance):
  """Checks whether the convergence criteria have been met."""
  # Check func tolerance
  value_range = tf.math.abs(
      tf.math.reduce_max(input_tensor=population_values) -
      tf.math.reduce_min(input_tensor=population_values))
  value_converged = value_range <= func_tolerance
  # Ideally, we would compute the position convergence by computing the
  # pairwise distance between every member of the population and checking if
  # the maximum of those is less than the supplied tolerance. However, this is
  # completely infeasible in terms of performance. We adopt a more conservative
  # approach which checks the distance between the first population member
  # with the rest of the population. If the largest such distance is less than
  # half the supplied tolerance, we stop. The reason why this is sufficient is
  # as follows. For any pair of distinct points (a, b) in the population, we
  # have the relation:  |a - b| <= |x0 - a| + |x0 - b|, where x0 is any
  # other point. In particular, let x0 be the first element of the population
  # and suppose that the largest distance between this point and any other
  # member is epsilon. Then, for any pair of points (a, b),
  # |a - b| <= 2 * epsilon and hence, the maximum distance between any pair of
  # points in the population is bounded above by twice the distance between
  # the first point and other points.
  half_tol = position_tolerance / 2
  def part_converged(part):
    return tf.math.reduce_max(input_tensor=tf.math.abs(part -
                                                       part[0])) <= half_tol

  x_converged = tf.math.reduce_all(
      input_tensor=[part_converged(part) for part in population])
  return value_converged | x_converged

def _get_starting_population(initial_population,
                             initial_position,
                             population_size,
                             population_stddev,
                             seed):
  """Constructs the initial population.

  If an initial population is not already provided, this function constructs
  a population by adding random normal noise to the initial position.

  Args:
    initial_population: None or a list of `Tensor`s. The initial population.
    initial_position: None or a list of `Tensor`s. The initial position.
      If initial_population is None, this argument must not be None.
    population_size: Scalar integer `Tensor`. The number of members in the
      population. If the initial population is not None, this parameter is
      ignored.
    population_stddev: A positive scalar real `Tensor` of the same dtype
      as `initial_position` or `initial_population` (whichever is not None).
      This parameter is ignored if `initial_population`
      is specified. Used to generate the population from the
      `initial_position` by adding random normal noise with zero mean and
      the specified standard deviation.
    seed: Seed for random number generation.

  Returns:
    A list of `Tensor`s. The initial population.
  """
  if initial_population is not None:
    return [tf.convert_to_tensor(value=part) for part in initial_population]
  # Constructs the population by adding normal noise to the initial position.
  seed_stream = distributions.SeedStream(seed, salt='get_starting_population')
  population = []
  for part in initial_position:
    part = tf.convert_to_tensor(value=part)
    part_event_shape = tf.shape(input=part)
    # We only draw population_size-1 random vectors because we want to ensure
    # that the supplied position is part of the population. The first member
    # is set to be the initial_position.
    population_part_shape = tf.concat([[population_size-1],
                                       part_event_shape], axis=0)
    population_part = tf.random.normal(population_part_shape,
                                       stddev=population_stddev,
                                       dtype=part.dtype.base_dtype,
                                       seed=seed_stream())
    population_part += part
    population_part = tf.concat([[part], population_part], axis=0)
    population.append(population_part)
  return population

def _binary_crossover(population,
                      population_size,
                      mutants,
                      crossover_prob,
                      seed):
  """Performs recombination by binary crossover for the current population.

  Let v_i denote the i'th component of the member v and m_i the corresponding
  component of the mutant vector corresponding to v. Then the crossed over
  vector w_i is determined by setting w_i =
  (m_i with probability=crossover_prob else v_i). In addition, DE requires that
  at least one of the components is crossed over (otherwise we end
  up with no change). This is done by choosing on index say k randomly where
  a force crossover is performed (i.e. w_k = m_k). This is the scheme
  implemented in this function.

  Args:
    population: A Python list of `Tensor`s where each `Tensor` in the list
      must be of rank at least 1 and all the elements must have a common
      first dimension. The base population to cross over.
    population_size: A scalar integer `Tensor`. The number of elements in the
      population (i.e. size of the first dimension of any member of
      `population`).
    mutants: A Python list of `Tensor`s with the same structure as `population`.
      The mutated population.
    crossover_prob: A postive real scalar `Tensor` bounded above by 1.0. The
      probability of a crossover being performed for each axis.
    seed: `int` or None. The random seed for this `Op`. If `None`, no seed is
      applied.

  Returns:
    A list of `Tensor`s of the same structure, dtype and shape as `population`.
    The recombined population.
  """
  sizes = [tf.cast(tf.size(input=x), dtype=tf.float64) for x in population]
  seed_stream = distributions.SeedStream(seed, salt='binary_crossover')
  force_crossover_group = distributions.Categorical(sizes).sample(
      [population_size, 1], seed=seed_stream())
  recombinants = []
  for i, population_part in enumerate(population):
    pop_part_flat = tf.reshape(population_part, [population_size, -1])
    mutant_part_flat = tf.reshape(mutants[i], [population_size, -1])
    part_size = tf.size(input=population_part) // population_size
    force_crossovers = tf.one_hot(
        tf.random.uniform([population_size],
                          minval=0,
                          maxval=part_size,
                          dtype=tf.int32,
                          seed=seed_stream()),
        part_size,
        on_value=True,
        off_value=False,
        dtype=tf.bool)  # Tensor of shape [population_size, size]
    group_mask = tf.math.equal(force_crossover_group, i)
    force_crossovers &= group_mask
    do_binary_crossover = tf.random.uniform(
        [population_size, part_size],
        dtype=crossover_prob.dtype.base_dtype,
        seed=seed_stream()) < crossover_prob
    do_binary_crossover |= force_crossovers
    recombinant_flat = tf.where(
        do_binary_crossover,
        x=mutant_part_flat,
        y=pop_part_flat)
    recombinant = tf.reshape(recombinant_flat, tf.shape(input=population_part))
    recombinants.append(recombinant)
  return recombinants

def _get_mutants(population,
                 population_size,
                 mixing_indices,
                 differential_weight):
  """Computes the mutatated vectors for each population member.

  Args:
    population:  Python `list` of `Tensor`s representing the
      current population vectors. Each `Tensor` must be of the same real dtype.
      The first dimension of each `Tensor` indexes individual
      population members. For example, if the population is a list with a
      single `Tensor` of shape [n, m1, m2], then `n` is the population size and
      the shape of an individual solution is [m1, m2].
      If there is more than one element in the population, then each `Tensor`
      in the list should have the first axis of the same size.
    population_size: Scalar integer `Tensor`. The size of the population.
    mixing_indices: `Tensor` of integral dtype and shape [n, 3] where `n` is the
      number of members in the population. Each element of the `Tensor` must be
      a valid index into the first dimension of the population (i.e range
      between `0` and `n-1` inclusive).
    differential_weight: Real scalar `Tensor`. Must be positive and less than
      2.0. The parameter controlling the strength of mutation.

  Returns:
    mutants: `Tensor` or Python `list` of `Tensor`s of the same shape and dtype
      as the input population. The mutated vectors.
  """
  mixing_indices = tf.reshape(mixing_indices, [-1])
  weights = tf.stack([1.0, differential_weight, -differential_weight])
  def _mutant_part(population_part):
    donors = tf.gather(population_part, mixing_indices)
    donors = tf.transpose(
        a=tf.reshape(donors, [population_size, 3, -1]), perm=[0, 2, 1])
    return tf.math.reduce_sum(input_tensor=donors * weights, axis=-1)

  return [_mutant_part(population_part) for population_part in population]

def _get_mixing_indices(size, seed=None, name=None):
  """Generates an array of indices suitable for mutation operation.

  The mutation operation in differential evolution requires that for every
  element of the population, three distinct other elements be chosen to produce
  a trial candidate. This function generates an array of shape [size, 3]
  satisfying the properties that:
    (a). array[i, :] does not contain the index 'i'.
    (b). array[i, :] does not contain any overlapping indices.
    (c). All elements in the array are between 0 and size - 1 inclusive.

  Args:
    size: Scalar integer `Tensor`. The number of samples as well as a the range
      of the indices to sample from.
    seed: `int` or None. The random seed for this `Op`. If `None`, no seed is
      applied.
      Default value: `None`.
    name:  Python `str` name prefixed to Ops created by this function.
      Default value: 'get_mixing_indices'.

  Returns:
    sample: A `Tensor` of shape [size, 3] and same dtype as `size` containing
    samples without replacement between 0 and size - 1 (inclusive) with the
    `i`th row not including the number `i`.
  """
  with tf.compat.v1.name_scope(
      name, default_name='get_mixing_indices', values=[size]):
    size = tf.convert_to_tensor(value=size)
    dtype = size.dtype
    seed_stream = distributions.SeedStream(seed, salt='get_mixing_indices')
    first = tf.random.uniform([size],
                              maxval=size-1,
                              dtype=dtype,
                              seed=seed_stream())
    second = tf.random.uniform([size],
                               maxval=size-2,
                               dtype=dtype,
                               seed=seed_stream())
    third = tf.random.uniform([size],
                              maxval=size-3,
                              dtype=dtype,
                              seed=seed_stream())

    # Shift second if it is on top of or to the right of first
    second = tf.where(first < second, x=second, y=second + 1)
    smaller = tf.math.minimum(first, second)
    larger = tf.math.maximum(first, second)
    # Shift the third one so it does not coincide with either the first or the
    # second number. Assuming first < second, shift by 1 if the number is in
    # [first, second) and by 2 if the number is greater than or equal to the
    # second.
    third = tf.where(third < smaller, x=third, y=third + 1)
    third = tf.where(third < larger, x=third, y=third + 1)
    sample = tf.stack([first, second, third], axis=1)
    to_avoid = tf.expand_dims(tf.range(size), axis=-1)
    sample = tf.where(sample < to_avoid, x=sample, y=sample + 1)
    return sample

def _ensure_list(tensor_or_list):
  """Converts the input arg to a list if it is not a list already.

  Args:
    tensor_or_list: A `Tensor` or a Python list of `Tensor`s. The argument to
      convert to a list of `Tensor`s.

  Returns:
    A tuple of two elements. The first is a Python list of `Tensor`s containing
    the original arguments. The second is a boolean indicating whether
    the original argument was a list or tuple already.
  """
  if isinstance(tensor_or_list, (list, tuple)):
    return list(tensor_or_list), True
  return [tensor_or_list], False


# Copyright 2018 The TensorFlow Probability Authors.
#
# 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.
# ============================================================================
"""The Deterministic distribution class."""

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

import abc

# Dependency imports
import six
import tensorflow.compat.v2 as tf

from tensorflow_probability.python.distributions import distribution
from tensorflow_probability.python.distributions import kullback_leibler
from tensorflow_probability.python.internal import assert_util
from tensorflow_probability.python.internal import distribution_util
from tensorflow_probability.python.internal import dtype_util
from tensorflow_probability.python.internal import reparameterization
from tensorflow_probability.python.internal import tensorshape_util

__all__ = [
    "Deterministic",
    "VectorDeterministic",
]

def _get_tol(tol, dtype, validate_args):
  """Gets a Tensor of type `dtype`, 0 if `tol` is None, validation optional."""
  if tol is None:
    return tf.convert_to_tensor(value=0, dtype=dtype)

  tol = tf.convert_to_tensor(value=tol, dtype=dtype)
  if validate_args:
    tol = distribution_util.with_dependencies([
        assert_util.assert_non_negative(
            tol, message="Argument 'tol' must be non-negative")
    ], tol)
  return tol

@six.add_metaclass(abc.ABCMeta)
class _BaseDeterministic(distribution.Distribution):
  """Base class for Deterministic distributions."""

  def __init__(self,
               loc,
               atol=None,
               rtol=None,
               is_vector=False,
               validate_args=False,
               allow_nan_stats=True,
               parameters=None,
               name="_BaseDeterministic"):
    """Initialize a batch of `_BaseDeterministic` distributions.

    The `atol` and `rtol` parameters allow for some slack in `pmf`, `cdf`
    computations, e.g. due to floating-point error.

    ```
    pmf(x; loc)
      = 1, if Abs(x - loc) <= atol + rtol * Abs(loc),
      = 0, otherwise.
    ```

    Args:
      loc: Numeric `Tensor`.  The point (or batch of points) on which this
        distribution is supported.
      atol:  Non-negative `Tensor` of same `dtype` as `loc` and broadcastable
        shape.  The absolute tolerance for comparing closeness to `loc`.
        Default is `0`.
      rtol:  Non-negative `Tensor` of same `dtype` as `loc` and broadcastable
        shape.  The relative tolerance for comparing closeness to `loc`.
        Default is `0`.
      is_vector:  Python `bool`.  If `True`, this is for `VectorDeterministic`,
        else `Deterministic`.
      validate_args: Python `bool`, default `False`. When `True` distribution
        parameters are checked for validity despite possibly degrading runtime
        performance. When `False` invalid inputs may silently render incorrect
        outputs.
      allow_nan_stats: Python `bool`, default `True`. When `True`, statistics
        (e.g., mean, mode, variance) use the value "`NaN`" to indicate the
        result is undefined. When `False`, an exception is raised if one or
        more of the statistic's batch members are undefined.
      parameters: Dict of locals to facilitate copy construction.
      name: Python `str` name prefixed to Ops created by this class.

    Raises:
      ValueError:  If `loc` is a scalar.
    """
    with tf.name_scope(name) as name:
      dtype = dtype_util.common_dtype([loc, atol, rtol],
                                      preferred_dtype=tf.float32)
      loc = tf.convert_to_tensor(value=loc, name="loc", dtype=dtype)
      if is_vector and validate_args:
        msg = "Argument loc must be at least rank 1."
        if tensorshape_util.rank(loc.shape) is not None:
          if tensorshape_util.rank(loc.shape) < 1:
            raise ValueError(msg)
        else:
          loc = distribution_util.with_dependencies(
              [assert_util.assert_rank_at_least(loc, 1, message=msg)], loc)
      self._loc = loc
      self._atol = _get_tol(atol, self._loc.dtype, validate_args)
      self._rtol = _get_tol(rtol, self._loc.dtype, validate_args)

      super(_BaseDeterministic, self).__init__(
          dtype=self._loc.dtype,
          reparameterization_type=reparameterization.NOT_REPARAMETERIZED,
          validate_args=validate_args,
          allow_nan_stats=allow_nan_stats,
          parameters=parameters,
          graph_parents=[self._loc, self._atol, self._rtol],
          name=name)

      # Avoid using the large broadcast with self.loc if possible.
      if rtol is None:
        self._slack = self.atol
      else:
        self._slack = self.atol + self.rtol * tf.abs(self.loc)

  @property
  def loc(self):
    """Point (or batch of points) at which this distribution is supported."""
    return self._loc

  @property
  def atol(self):
    """Absolute tolerance for comparing points to `self.loc`."""
    return self._atol

  @property
  def rtol(self):
    """Relative tolerance for comparing points to `self.loc`."""
    return self._rtol

  def _entropy(self):
    return tf.zeros(self.batch_shape_tensor(), dtype=self.dtype)

  def _mean(self):
    return tf.identity(self.loc)

  def _variance(self):
    return tf.zeros_like(self.loc)

  def _mode(self):
    return self.mean()

  def _sample_n(self, n, seed=None):
    del seed  # unused
    return tf.broadcast_to(
        self.loc,
        tf.concat([[n], self.batch_shape_tensor(), self.event_shape_tensor()],
                  axis=0))

class Deterministic(_BaseDeterministic):
  """Scalar `Deterministic` distribution on the real line.

  The scalar `Deterministic` distribution is parameterized by a [batch] point
  `loc` on the real line.  The distribution is supported at this point only,
  and corresponds to a random variable that is constant, equal to `loc`.

  See [Degenerate rv](https://en.wikipedia.org/wiki/Degenerate_distribution).

  #### Mathematical Details

  The probability mass function (pmf) and cumulative distribution function (cdf)
  are

  ```none
  pmf(x; loc) = 1, if x == loc, else 0
  cdf(x; loc) = 1, if x >= loc, else 0
  ```

  #### Examples

  ```python
  # Initialize a single Deterministic supported at zero.
  constant = tfp.distributions.Deterministic(0.)
  constant.prob(0.)
  ==> 1.
  constant.prob(2.)
  ==> 0.

  # Initialize a [2, 2] batch of scalar constants.
  loc = [[0., 1.], [2., 3.]]
  x = [[0., 1.1], [1.99, 3.]]
  constant = tfp.distributions.Deterministic(loc)
  constant.prob(x)
  ==> [[1., 0.], [0., 1.]]
  ```

  """

  def __init__(self,
               loc,
               atol=None,
               rtol=None,
               validate_args=False,
               allow_nan_stats=True,
               name="Deterministic"):
    """Initialize a scalar `Deterministic` distribution.

    The `atol` and `rtol` parameters allow for some slack in `pmf`, `cdf`
    computations, e.g. due to floating-point error.

    ```
    pmf(x; loc)
      = 1, if Abs(x - loc) <= atol + rtol * Abs(loc),
      = 0, otherwise.
    ```

    Args:
      loc: Numeric `Tensor` of shape `[B1, ..., Bb]`, with `b >= 0`.
        The point (or batch of points) on which this distribution is supported.
      atol:  Non-negative `Tensor` of same `dtype` as `loc` and broadcastable
        shape.  The absolute tolerance for comparing closeness to `loc`.
        Default is `0`.
      rtol:  Non-negative `Tensor` of same `dtype` as `loc` and broadcastable
        shape.  The relative tolerance for comparing closeness to `loc`.
        Default is `0`.
      validate_args: Python `bool`, default `False`. When `True` distribution
        parameters are checked for validity despite possibly degrading runtime
        performance. When `False` invalid inputs may silently render incorrect
        outputs.
      allow_nan_stats: Python `bool`, default `True`. When `True`, statistics
        (e.g., mean, mode, variance) use the value "`NaN`" to indicate the
        result is undefined. When `False`, an exception is raised if one or
        more of the statistic's batch members are undefined.
      name: Python `str` name prefixed to Ops created by this class.
    """
    parameters = dict(locals())
    super(Deterministic, self).__init__(
        loc,
        atol=atol,
        rtol=rtol,
        validate_args=validate_args,
        allow_nan_stats=allow_nan_stats,
        parameters=parameters,
        name=name)

  @classmethod
  def _params_event_ndims(cls):
    return dict(loc=0, atol=0, rtol=0)

  def _batch_shape_tensor(self):
    return tf.broadcast_dynamic_shape(
        tf.shape(input=self.loc), tf.shape(input=self._slack))

  def _batch_shape(self):
    return tf.broadcast_static_shape(self.loc.shape, self._slack.shape)

  def _event_shape_tensor(self):
    return tf.constant([], dtype=tf.int32)

  def _event_shape(self):
    return tf.TensorShape([])

  def _prob(self, x):
    return tf.cast(tf.abs(x - self.loc) <= self._slack, dtype=self.dtype)

  def _cdf(self, x):
    return tf.cast(x >= self.loc - self._slack, dtype=self.dtype)

class VectorDeterministic(_BaseDeterministic):
  """Vector `Deterministic` distribution on `R^k`.

  The `VectorDeterministic` distribution is parameterized by a [batch] point
  `loc in R^k`.  The distribution is supported at this point only,
  and corresponds to a random variable that is constant, equal to `loc`.

  See [Degenerate rv](https://en.wikipedia.org/wiki/Degenerate_distribution).

  #### Mathematical Details

  The probability mass function (pmf) is

  ```none
  pmf(x; loc)
    = 1, if All[Abs(x - loc) <= atol + rtol * Abs(loc)],
    = 0, otherwise.
  ```

  #### Examples

  ```python
  tfd = tfp.distributions

  # Initialize a single VectorDeterministic supported at [0., 2.] in R^2.
  constant = tfd.Deterministic([0., 2.])
  constant.prob([0., 2.])
  ==> 1.
  constant.prob([0., 3.])
  ==> 0.

  # Initialize a [3] batch of constants on R^2.
  loc = [[0., 1.], [2., 3.], [4., 5.]]
  constant = tfd.VectorDeterministic(loc)
  constant.prob([[0., 1.], [1.9, 3.], [3.99, 5.]])
  ==> [1., 0., 0.]
  ```

  """

  def __init__(self,
               loc,
               atol=None,
               rtol=None,
               validate_args=False,
               allow_nan_stats=True,
               name="VectorDeterministic"):
    """Initialize a `VectorDeterministic` distribution on `R^k`, for `k >= 0`.

    Note that there is only one point in `R^0`, the "point" `[]`.  So if `k = 0`
    then `self.prob([]) == 1`.

    The `atol` and `rtol` parameters allow for some slack in `pmf`
    computations, e.g. due to floating-point error.

    ```
    pmf(x; loc)
      = 1, if All[Abs(x - loc) <= atol + rtol * Abs(loc)],
      = 0, otherwise
    ```

    Args:
      loc: Numeric `Tensor` of shape `[B1, ..., Bb, k]`, with `b >= 0`, `k >= 0`
        The point (or batch of points) on which this distribution is supported.
      atol:  Non-negative `Tensor` of same `dtype` as `loc` and broadcastable
        shape.  The absolute tolerance for comparing closeness to `loc`.
        Default is `0`.
      rtol:  Non-negative `Tensor` of same `dtype` as `loc` and broadcastable
        shape.  The relative tolerance for comparing closeness to `loc`.
        Default is `0`.
      validate_args: Python `bool`, default `False`. When `True` distribution
        parameters are checked for validity despite possibly degrading runtime
        performance. When `False` invalid inputs may silently render incorrect
        outputs.
      allow_nan_stats: Python `bool`, default `True`. When `True`, statistics
        (e.g., mean, mode, variance) use the value "`NaN`" to indicate the
        result is undefined. When `False`, an exception is raised if one or
        more of the statistic's batch members are undefined.
      name: Python `str` name prefixed to Ops created by this class.
    """
    parameters = dict(locals())
    super(VectorDeterministic, self).__init__(
        loc,
        atol=atol,
        rtol=rtol,
        is_vector=True,
        validate_args=validate_args,
        allow_nan_stats=allow_nan_stats,
        parameters=parameters,
        name=name)

  @classmethod
  def _params_event_ndims(cls):
    return dict(loc=1, atol=1, rtol=1)

  def _batch_shape_tensor(self):
    return tf.broadcast_dynamic_shape(
        tf.shape(input=self.loc), tf.shape(input=self._slack))[:-1]

  def _batch_shape(self):
    return tf.broadcast_static_shape(self.loc.shape, self._slack.shape)[:-1]

  def _event_shape_tensor(self):
    return tf.shape(input=self.loc)[-1:]

  def _event_shape(self):
    return self.loc.shape[-1:]

  def _prob(self, x):
    if self.validate_args:
      is_vector_check = assert_util.assert_rank_at_least(x, 1)
      right_vec_space_check = assert_util.assert_equal(
          self.event_shape_tensor(),
          tf.gather(tf.shape(input=x),
                    tf.rank(x) - 1),
          message="Argument 'x' not defined in the same space R^k as this distribution"
      )
      with tf.control_dependencies([is_vector_check]):
        with tf.control_dependencies([right_vec_space_check]):
          x = tf.identity(x)
    return tf.cast(
        tf.reduce_all(
            input_tensor=tf.abs(x - self.loc) <= self._slack, axis=-1),
        dtype=self.dtype)

@kullback_leibler.RegisterKL(_BaseDeterministic, distribution.Distribution)
def _kl_deterministic_distribution(a, b, name=None):
  """Calculate the batched KL divergence `KL(a || b)` with `a` Deterministic.

  Args:
    a: instance of a Deterministic distribution object.
    b: instance of a Distribution distribution object.
    name: (optional) Name to use for created operations. Default is
      "kl_deterministic_distribution".

  Returns:
    Batchwise `KL(a || b)`.
  """
  with tf.name_scope(name or "kl_deterministic_distribution"):
    return -b.log_prob(a.loc)


# Copyright 2018 The TensorFlow Probability Authors.
#
# 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.
# ============================================================================
"""SinhArcsinh bijector."""

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

# Dependency imports
import numpy as np

import tensorflow.compat.v2 as tf

from tensorflow_probability.python.bijectors import bijector
from tensorflow_probability.python.internal import assert_util
from tensorflow_probability.python.internal import distribution_util
from tensorflow_probability.python.internal import dtype_util

__all__ = [
    "SinhArcsinh",
]

def _sqrtx2p1(x):
  """Implementation of `sqrt(1 + x**2)` which is stable despite large `x`."""
  sqrt_eps = np.sqrt(np.finfo(dtype_util.as_numpy_dtype(x.dtype)).eps)
  return tf.where(
      tf.abs(x) * sqrt_eps <= 1.,
      tf.sqrt(x**2. + 1.),
      # For large x, calculating x**2 can overflow. This can be alleviated by
      # considering:
      # sqrt(1 + x**2)
      # = exp(0.5 log(1 + x**2))
      # = exp(0.5 log(x**2 * (1 + x**-2)))
      # = exp(log(x) + 0.5 * log(1 + x**-2))
      # = |x| * exp(0.5 log(1 + x**-2))
      # = |x| * sqrt(1 + x**-2)
      # We omit the last term in this approximation.
      # When |x| > 1 / sqrt(machineepsilon), the second term will be 1,
      # due to sqrt(1 + x**-2) = 1. This is also true with the gradient term,
      # and higher order gradients, since the first order derivative of
      # sqrt(1 + x**-2) is -2 * x**-3 / (1 + x**-2) = -2 / (x**3 + x),
      # and all nth-order derivatives will be O(x**-(n + 2)). This makes any
      # gradient terms that contain any derivatives of sqrt(1 + x**-2) vanish.
      tf.abs(x))

class SinhArcsinh(bijector.Bijector):
  """Compute `Y = g(X) = Sinh( (Arcsinh(X) + skewness) * tailweight )`.

  For `skewness in (-inf, inf)` and `tailweight in (0, inf)`, this
  transformation is a
  diffeomorphism of the real line `(-inf, inf)`.  The inverse transform is
  `X = g^{-1}(Y) = Sinh( ArcSinh(Y) / tailweight - skewness )`.

  The `SinhArcsinh` transformation of the Normal is described in
  [Sinh-arcsinh distributions](https://www.jstor.org/stable/27798865)
  This Bijector allows a similar transformation of any distribution supported on
  `(-inf, inf)`.

  #### Meaning of the parameters

  * If `skewness = 0` and `tailweight = 1`, this transform is the identity.
  * Positive (negative) `skewness` leads to positive (negative) skew.
    * positive skew means, for unimodal `X` centered at zero, the mode of `Y` is
      "tilted" to the right.
    * positive skew means positive values of `Y` become more likely, and
      negative values become less likely.
  * Larger (smaller) `tailweight` leads to fatter (thinner) tails.
    * Fatter tails mean larger values of `|Y|` become more likely.
    * If `X` is a unit Normal, `tailweight < 1` leads to a distribution that is
      "flat" around `Y = 0`, and a very steep drop-off in the tails.
    * If `X` is a unit Normal, `tailweight > 1` leads to a distribution more
      peaked at the mode with heavier tails.

  To see the argument about the tails, note that for `|X| >> 1` and
  `|X| >> (|skewness| * tailweight)**tailweight`, we have
  `Y approx 0.5 X**tailweight e**(sign(X) skewness * tailweight)`.
  """

  def __init__(self,
               skewness=None,
               tailweight=None,
               validate_args=False,
               name="SinhArcsinh"):
    """Instantiates the `SinhArcsinh` bijector.

    Args:
      skewness:  Skewness parameter.  Float-type `Tensor`.  Default is `0`
        of type `float32`.
      tailweight:  Tailweight parameter.  Positive `Tensor` of same `dtype` as
        `skewness` and broadcastable `shape`.  Default is `1` of type `float32`.
      validate_args: Python `bool` indicating whether arguments should be
        checked for correctness.
      name: Python `str` name given to ops managed by this object.
    """
    self._graph_parents = []
    self._name = name
    self._validate_args = validate_args
    with self._name_scope("init"):
      tailweight = 1. if tailweight is None else tailweight
      skewness = 0. if skewness is None else skewness
      self._skewness = tf.convert_to_tensor(value=skewness, name="skewness")
      self._tailweight = tf.convert_to_tensor(
          value=tailweight, name="tailweight", dtype=self._skewness.dtype)
      dtype_util.assert_same_float_dtype([self._skewness, self._tailweight])
      if validate_args:
        self._tailweight = distribution_util.with_dependencies([
            assert_util.assert_positive(
                self._tailweight,
                message="Argument tailweight was not positive")
        ], self._tailweight)
    super(SinhArcsinh, self).__init__(
        forward_min_event_ndims=0,
        validate_args=validate_args,
        name=name)

  @property
  def skewness(self):
    """The `skewness` in: `Y  = Sinh((Arcsinh(X) + skewness) * tailweight)`."""
    return self._skewness

  @property
  def tailweight(self):
    """The `tailweight` in: `Y = Sinh((Arcsinh(X) + skewness) * tailweight)`."""
    return self._tailweight

  def _forward(self, x):
    return tf.sinh((tf.asinh(x) + self.skewness) * self.tailweight)

  def _inverse(self, y):
    return tf.sinh(tf.asinh(y) / self.tailweight - self.skewness)

  def _inverse_log_det_jacobian(self, y):
    # x = sinh(arcsinh(y) / tailweight - skewness)
    # Using sinh' = cosh, arcsinh'(y) = 1 / sqrt(y**2 + 1),
    # dx/dy
    # = cosh(arcsinh(y) / tailweight - skewness)
    #     / (tailweight * sqrt(y**2 + 1))

    # This is computed inside the log to avoid catastrophic cancellations
    # from cosh((arcsinh(y) / tailweight) - skewness) and sqrt(x**2 + 1).
    return (tf.math.log(
        tf.cosh(tf.asinh(y) / self.tailweight - self.skewness)
        # TODO(srvasude): Consider using cosh(arcsinh(x)) in cases
        # where (arcsinh(x) / tailweight) - skewness ~= arcsinh(x).
        / _sqrtx2p1(y)) - tf.math.log(self.tailweight))

  def _forward_log_det_jacobian(self, x):
    # y = sinh((arcsinh(x) + skewness) * tailweight)
    # Using sinh' = cosh, arcsinh'(x) = 1 / sqrt(x**2 + 1),
    # dy/dx
    # = cosh((arcsinh(x) + skewness) * tailweight) * tailweight / sqrt(x**2 + 1)

    # This is computed inside the log to avoid catastrophic cancellations
    # from cosh((arcsinh(x) + skewness) * tailweight) and sqrt(x**2 + 1).
    return (tf.math.log(
        tf.cosh((tf.asinh(x) + self.skewness) * self.tailweight)
        # TODO(srvasude): Consider using cosh(arcsinh(x)) in cases
        # where (arcsinh(x) + skewness) * tailweight ~= arcsinh(x).
        / _sqrtx2p1(x)) + tf.math.log(self.tailweight))


# Copyright 2018 The TensorFlow Probability Authors.
#
# 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.
# ============================================================================
"""Numerically stable variants of common mathematical expressions."""

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

import numpy as np
import tensorflow as tf

def log1psquare(x, name=None):
  """Numerically stable calculation of `log(1 + x**2)` for small or large `|x|`.

  For sufficiently large `x` we use the following observation:

  ```none
  log(1 + x**2) =   2 log(|x|) + log(1 + 1 / x**2)
                --> 2 log(|x|)  as x --> inf
  ```

  Numerically, `log(1 + 1 / x**2)` is `0` when `1 / x**2` is small relative to
  machine epsilon.

  Args:
    x: Float `Tensor` input.
    name: Python string indicating the name of the TensorFlow operation.
      Default value: `'log1psquare'`.

  Returns:
    log1psq: Float `Tensor` representing `log(1. + x**2.)`.
  """
  with tf.compat.v1.name_scope(name, 'log1psquare', [x]):
    x = tf.convert_to_tensor(value=x, dtype_hint=tf.float32, name='x')
    dtype = x.dtype.as_numpy_dtype

    eps = np.finfo(dtype).eps.astype(np.float64)
    is_large = tf.abs(x) > (eps**-0.5).astype(dtype)

    # Mask out small x's so the gradient correctly propagates.
    abs_large_x = tf.where(is_large, tf.abs(x), tf.ones_like(x))
    return tf.where(is_large, 2. * tf.math.log(abs_large_x),
                    tf.math.log1p(tf.square(x)))

def soft_threshold(x, threshold, name=None):
  """Soft Thresholding operator.

  This operator is defined by the equations

  ```none
                                { x[i] - gamma,  x[i] >   gamma
  SoftThreshold(x, gamma)[i] =  { 0,             x[i] ==  gamma
                                { x[i] + gamma,  x[i] <  -gamma
  ```

  In the context of proximal gradient methods, we have

  ```none
  SoftThreshold(x, gamma) = prox_{gamma L1}(x)
  ```

  where `prox` is the proximity operator.  Thus the soft thresholding operator
  is used in proximal gradient descent for optimizing a smooth function with
  (non-smooth) L1 regularization, as outlined below.

  The proximity operator is defined as:

  ```none
  prox_r(x) = argmin{ r(z) + 0.5 ||x - z||_2**2 : z },
  ```

  where `r` is a (weakly) convex function, not necessarily differentiable.
  Because the L2 norm is strictly convex, the above argmin is unique.

  One important application of the proximity operator is as follows.  Let `L` be
  a convex and differentiable function with Lipschitz-continuous gradient.  Let
  `R` be a convex lower semicontinuous function which is possibly
  nondifferentiable.  Let `gamma` be an arbitrary positive real.  Then

  ```none
  x_star = argmin{ L(x) + R(x) : x }
  ```

  if and only if the fixed-point equation is satisfied:

  ```none
  x_star = prox_{gamma R}(x_star - gamma grad L(x_star))
  ```

  Proximal gradient descent thus typically consists of choosing an initial value
  `x^{(0)}` and repeatedly applying the update

  ```none
  x^{(k+1)} = prox_{gamma^{(k)} R}(x^{(k)} - gamma^{(k)} grad L(x^{(k)}))
  ```

  where `gamma` is allowed to vary from iteration to iteration.  Specializing to
  the case where `R(x) = ||x||_1`, we minimize `L(x) + ||x||_1` by repeatedly
  applying the update

  ```
  x^{(k+1)} = SoftThreshold(x - gamma grad L(x^{(k)}), gamma)
  ```

  (This idea can also be extended to second-order approximations, although the
  multivariate case does not have a known closed form like above.)

  Args:
    x: `float` `Tensor` representing the input to the SoftThreshold function.
    threshold: nonnegative scalar, `float` `Tensor` representing the radius of
      the interval on which each coordinate of SoftThreshold takes the value
      zero.  Denoted `gamma` above.
    name: Python string indicating the name of the TensorFlow operation.
      Default value: `'soft_threshold'`.

  Returns:
    softthreshold: `float` `Tensor` with the same shape and dtype as `x`,
      representing the value of the SoftThreshold function.

  #### References

  [1]: Yu, Yao-Liang. The Proximity Operator.
       https://www.cs.cmu.edu/~suvrit/teach/yaoliang_proximity.pdf

  [2]: Wikipedia Contributors. Proximal gradient methods for learning.
       _Wikipedia, The Free Encyclopedia_, 2018.
       https://en.wikipedia.org/wiki/Proximal_gradient_methods_for_learning

  """
  # https://math.stackexchange.com/questions/471339/derivation-of-soft-thresholding-operator
  with tf.compat.v1.name_scope(name, 'soft_threshold', [x, threshold]):
    x = tf.convert_to_tensor(value=x, name='x')
    threshold = tf.convert_to_tensor(
        value=threshold, dtype=x.dtype, name='threshold')
    return tf.sign(x) * tf.maximum(tf.abs(x) - threshold, 0.)

def clip_by_value_preserve_gradient(t, clip_value_min, clip_value_max,
                                    name=None):
  """Clips values to a specified min and max while leaving gradient unaltered.

  Like `tf.clip_by_value`, this function returns a tensor of the same type and
  shape as input `t` but with values clamped to be no smaller than to
  `clip_value_min` and no larger than `clip_value_max`. Unlike
  `tf.clip_by_value`, the gradient is unaffected by this op, i.e.,

  ```python
  tf.gradients(tfp.math.clip_by_value_preserve_gradient(x), x)[0]
  # ==> ones_like(x)
  ```

  Note: `clip_value_min` needs to be smaller or equal to `clip_value_max` for
  correct results.

  Args:
    t: A `Tensor`.
    clip_value_min: A scalar `Tensor`, or a `Tensor` with the same shape
      as `t`. The minimum value to clip by.
    clip_value_max: A scalar `Tensor`, or a `Tensor` with the same shape
      as `t`. The maximum value to clip by.
    name: A name for the operation (optional).
      Default value: `'clip_by_value_preserve_gradient'`.

  Returns:
    clipped_t: A clipped `Tensor`.
  """
  with tf.compat.v1.name_scope(name, 'clip_by_value_preserve_gradient',
                               [t, clip_value_min, clip_value_max]):
    t = tf.convert_to_tensor(value=t, name='t')
    clip_t = tf.clip_by_value(t, clip_value_min, clip_value_max)
    return t + tf.stop_gradient(clip_t - t)


# Copyright 2018 The TensorFlow Probability Authors.
#
# 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.
# ============================================================================
"""Trains a GAN to generate synthetic images of MNIST digits.

A Generative Adversarial Network (GAN) is a generative model that learns the
probability distribution of training examples and generates similar samples. It
simultaneously trains two models: a generative model G that learns the data
distribution, and a discriminative model D that estimates the probability that a
sample came from the training data [1].

To learn the generator's distribution p_g over data z, we define a multilayer
perceptron G(z; theta_g), where the prior of input noise z~p_z(z) is a Gaussian
distribution. We also define a second multilayer perceptron D(x; theta_d)
that outputs a single scalar. D(x) represents the probability that x came
from the training data rather than p_g.

We train D to maximize the probability of assigning the correct label to both
training examples and samples from G. We simultaneously train G to maximize the
probability of D making a mistake. This framework corresponds to a minimax
two-player game with value function V(G, D):

```none
V(G, D) = E_{x~p_data(x)}[log(D(x))] + E_{z~p_z(z)}[log(1-D(G(z)))]
```

This optimization problem is bilevel: it requires a minima solution with respect
to generative parameters theta_g and a maxima solution with respect to
discriminative parameters theta_d. In practice, the algorithm proceeds by
iterating gradient updates on each network. The goal of training is to reach
the equilibrium where the generator produces samples that are indistinguishable
by the discriminator.

#### References

[1]: Ian J. Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu,
     David Warde-Farley, Sherjil Ozair, Aaron Courville, Yoshua Bengio.
     Generative Adversarial Nets. In _Neural Information Processing
     Systems Conference_, 2014.
     https://arxiv.org/pdf/1406.2661.pdf
"""

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

import os

# Dependency imports
from absl import flags
from matplotlib.backends import backend_agg
import matplotlib.pyplot as plt
import numpy as np
import tensorflow as tf
import tensorflow_probability as tfp

from tensorflow.contrib.learn.python.learn.datasets import mnist

tfd = tfp.distributions

IMAGE_SHAPE = [28, 28, 1]

tf.flags.DEFINE_float('learning_rate',
                      default=1e-4,
                      help='Initial learning rate.')
tf.flags.DEFINE_integer('max_steps',
                        default=100000,
                        help='Number of training steps to run.')
tf.flags.DEFINE_integer('batch_size',
                        default=128,
                        help='Batch size.')
tf.flags.DEFINE_integer('hidden_size',
                        default=128,
                        help='Hidden layer size.')
flags.DEFINE_integer('viz_steps',
                     default=1000,
                     help='Frequency at which save generated images.')
flags.DEFINE_string('data_dir',
                    default='/tmp/data',
                    help='Directory where data is stored (if using real data)')
flags.DEFINE_string(
    'model_dir',
    default=os.path.join(os.getenv('TEST_TMPDIR', '/tmp'),
                         'generative_adversarial_network/'),
    help="Directory to put the model's fit.")
flags.DEFINE_bool('fake_data',
                  default=None,
                  help='If true, uses fake data. Defaults to real data.')

FLAGS = flags.FLAGS

def build_input_pipeline(train_images, batch_size):
  """Build an iterator over training batches."""

  training_dataset = tf.data.Dataset.from_tensor_slices(train_images)
  training_batches = training_dataset.shuffle(
      50000, reshuffle_each_iteration=True).repeat().batch(batch_size)
  training_iterator = tf.compat.v1.data.make_one_shot_iterator(training_batches)
  images = training_iterator.get_next()
  return images

def build_fake_data(size):
  """Generate fake images of MNIST digits."""

  # Generate random noise from a Gaussian distribution.
  fake_images = np.random.normal(size=size)
  return fake_images

def plot_generated_images(images, fname):
  """Save a synthetic image as a PNG file.

  Args:
    images: samples of synthetic images generated by the generative network.
    fname: Python `str`, filename to save the plot to.
  """
  fig = plt.figure(figsize=(4, 4))
  canvas = backend_agg.FigureCanvasAgg(fig)

  for i, image in enumerate(images):
    ax = fig.add_subplot(4, 4, i + 1)
    plt.axis('off')
    ax.set_xticklabels([])
    ax.set_yticklabels([])
    ax.imshow(image.reshape(IMAGE_SHAPE[:-1]), cmap='Greys_r')

  fig.tight_layout()
  plt.subplots_adjust(wspace=0.05, hspace=0.05)
  canvas.print_figure(fname, format='png')

def main(argv):
  del argv  # unused
  if tf.io.gfile.exists(FLAGS.model_dir):
    tf.compat.v1.logging.warning(
        'Warning: deleting old log directory at {}'.format(FLAGS.model_dir))
    tf.io.gfile.rmtree(FLAGS.model_dir)
  tf.io.gfile.makedirs(FLAGS.model_dir)

  # Collapse the image data dimension for use with a fully-connected layer.
  image_size = np.prod(IMAGE_SHAPE, dtype=np.int32)
  if FLAGS.fake_data:
    train_images = build_fake_data([10, image_size])
  else:
    mnist_data = mnist.read_data_sets(FLAGS.data_dir, reshape=image_size)
    train_images = mnist_data.train.images

  images = build_input_pipeline(train_images, FLAGS.batch_size)

  # Build a Generative network. We use the Flipout Monte Carlo estimator
  # for the fully-connected layers: this enables lower variance stochastic
  # gradients than naive reparameterization.
  with tf.compat.v1.name_scope('Generator'):
    random_noise = tf.placeholder(tf.float64, shape=[None, FLAGS.hidden_size])
    generative_net = tf.keras.Sequential([
        tfp.layers.DenseFlipout(FLAGS.hidden_size, activation=tf.nn.relu),
        tfp.layers.DenseFlipout(image_size, activation=tf.sigmoid)
    ])
    sythetic_image = generative_net(random_noise)

  # Build a Discriminative network. Define the model as a Bernoulli
  # distribution parameterized by logits from a fully-connected layer.
  with tf.compat.v1.name_scope('Discriminator'):
    discriminative_net = tf.keras.Sequential([
        tfp.layers.DenseFlipout(FLAGS.hidden_size, activation=tf.nn.relu),
        tfp.layers.DenseFlipout(1)
    ])
    logits_real = discriminative_net(images)
    logits_fake = discriminative_net(sythetic_image)
    labels_distribution_real = tfd.Bernoulli(logits=logits_real)
    labels_distribution_fake = tfd.Bernoulli(logits=logits_fake)

  # Compute the model loss for discrimator and generator, averaged over
  # the batch size.
  loss_real = -tf.reduce_mean(
      input_tensor=labels_distribution_real.log_prob(
          tf.ones_like(logits_real)))
  loss_fake = -tf.reduce_mean(
      input_tensor=labels_distribution_fake.log_prob(
          tf.zeros_like(logits_fake)))
  loss_discriminator = loss_real + loss_fake
  loss_generator = -tf.reduce_mean(
      input_tensor=labels_distribution_fake.log_prob(
          tf.ones_like(logits_fake)))

  with tf.compat.v1.name_scope('train'):
    optimizer = tf.train.AdamOptimizer(learning_rate=FLAGS.learning_rate)
    train_op_discriminator = optimizer.minimize(
        loss_discriminator,
        var_list=tf.get_collection(
            tf.GraphKeys.TRAINABLE_VARIABLES, scope='Discriminator'))
    train_op_generator = optimizer.minimize(
        loss_generator,
        var_list=tf.get_collection(
            tf.GraphKeys.TRAINABLE_VARIABLES, scope='Generator'))

  with tf.compat.v1.Session() as sess:
    sess.run(tf.global_variables_initializer())
    for step in range(FLAGS.max_steps + 1):
      # Iterate gradient updates on each network.
      _, loss_value_d = sess.run([train_op_discriminator, loss_discriminator],
                                 feed_dict={random_noise: build_fake_data(
                                     [FLAGS.batch_size, FLAGS.hidden_size])})
      _, loss_value_g = sess.run([train_op_generator, loss_generator],
                                 feed_dict={random_noise: build_fake_data(
                                     [FLAGS.batch_size, FLAGS.hidden_size])})

      # Visualize some sythetic images produced by the generative network.
      if step % FLAGS.viz_steps == 0:
        images = sess.run(sythetic_image,
                          feed_dict={random_noise: build_fake_data(
                              [16, FLAGS.hidden_size])})

        plot_generated_images(images, fname=os.path.join(
            FLAGS.model_dir,
            'step{:06d}_images.png'.format(step)))

        print('Step: {:>3d} Loss_discriminator: {:.3f} '
              'Loss_generator: {:.3f}'.format(step, loss_value_d, loss_value_g))

if __name__ == '__main__':
  tf.compat.v1.app.run()


# Copyright 2018 The TensorFlow Probability Authors.
#
# 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.
# ============================================================================
"""Trains a grammar variational auto-encoder on synthetic data.

The grammar variational auto-encoder (VAE) [1] posits a generative model over
productions from a context-free grammar, and it posits an amortized variational
approximation for efficient posterior inference. We train the grammar VAE
on synthetic data using the grammar from [1] (Figure 1). Note for real data
analyses, one should implement a parser to convert examples into lists of
production rules.

This example showcases eager execution in order to train a model where data
points have a variable number of time steps (that is, without padding). However,
note that handling a variable number of time steps requires a batch size of 1.
In this example, we assume data points arrive in a stream, one at a time. Such a
setting has an unbounded maximum length which prevents padding.

Summaries are written under the flag `model_dir`. Point TensorBoard to that
directory in order to monitor progress.

Example output:

```none
Random examples from synthetic data distribution:
222N1N21c
1c2N2C2C12C1N
C11C12c
2C
NCC

Step:   0 Loss: -13.724 (0.494 sec)
Step: 500 Loss: -0.004 (145.741 sec)
Step: 1000 Loss: -0.000 (292.205 sec)
Step: 1500 Loss: -0.000 (438.819 sec)
```

#### References

[1]: Matt J. Kusner, Brooks Paige, and Jose Miguel Hernandez-Lobato. Grammar
     Variational Autoencoder. In _International Conference on Machine Learning_,
     2017. https://arxiv.org/abs/1703.01925
"""

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

import os
import time

# Dependency imports
from absl import flags
import six
import tensorflow as tf
import tensorflow_probability as tfp

from tensorflow_probability import edward2 as ed

flags.DEFINE_float("learning_rate",
                   default=1e-4,
                   help="Initial learning rate.")
flags.DEFINE_integer("max_steps",
                     default=5000,
                     help="Number of training steps to run.")
flags.DEFINE_integer("latent_size",
                     default=128,
                     help="Number of dimensions in the latent code.")
flags.DEFINE_integer("num_units",
                     default=256,
                     help="Number of units in the generative model's LSTM.")
flags.DEFINE_string("model_dir",
                    default=os.path.join(os.getenv("TEST_TMPDIR", "/tmp"),
                                         "grammar_vae/"),
                    help="Directory to put the model's fit.")

FLAGS = flags.FLAGS

class SmilesGrammar(object):
  """Context-free grammar for SMILES strings.

  A context-free grammar is a 4-tuple consisting of the following elements:

  + `nonterminal_symbols`: finite set of strings.
  + `alphabet`: finite set of strings (terminal symbols). It is disjoint from
      `nonterminal_symbols`.
  + `production_rules`: list of 2-tuples. The first and second elements of
      each tuple respectively denote the left-hand-side and right-hand-side of a
      production rule. All right-hand-sides are written as lists, since the
      number of right-hand-side symbols may be greater than 1.
  + `start_symbol`: string, a distinct nonterminal symbol.
  """

  @property
  def nonterminal_symbols(self):
    return {"smiles", "chain", "branched atom", "atom", "ringbond",
            "aromatic organic", "aliphatic organic", "digit"}

  @property
  def alphabet(self):
    return {"c", "C", "N", "1", "2"}

  @property
  def production_rules(self):
    return [
        ("smiles", ["chain"]),
        ("chain", ["chain", "branched atom"]),
        ("chain", ["branched atom"]),
        ("branched atom", ["atom", "ringbond"]),
        ("branched atom", ["atom"]),
        ("atom", ["aromatic organic"]),
        ("atom", ["aliphatic organic"]),
        ("ringbond", ["digit"]),
        ("aromatic organic", ["c"]),
        ("aliphatic organic", ["C"]),
        ("aliphatic organic", ["N"]),
        ("digit", ["1"]),
        ("digit", ["2"]),
    ]

  @property
  def start_symbol(self):
    return "smiles"

  def convert_to_string(self, productions):
    """Converts a sequence of productions into a string of terminal symbols.

    Args:
      productions: Tensor of shape [1, num_productions, num_production_rules].
        Slices along the `num_productions` dimension represent one-hot vectors.

    Returns:
      str that concatenates all terminal symbols from `productions`.

    Raises:
      ValueError: If the first production rule does not begin with
        `self.start_symbol`.
    """
    symbols = []
    for production in tf.unstack(productions, axis=1):
      lhs, rhs = self.production_rules[tf.argmax(input=production, axis=-1)]
      if not symbols:  # first iteration
        if lhs != self.start_symbol:
          raise ValueError("`productions` must begin with `self.start_symbol`.")
        symbols = rhs
      else:
        # Greedily unroll the nonterminal symbols based on the first occurrence
        # in a linear sequence.
        index = symbols.index(lhs)
        symbols = symbols[:index] + rhs + symbols[index + 1:]
    string = "".join(symbols)
    return string

  def mask(self, symbol, on_value, off_value):
    """Produces a masking tensor for (in)valid production rules.

    Args:
      symbol: str, a symbol in the grammar.
      on_value: Value to use for a valid production rule.
      off_value: Value to use for an invalid production rule.

    Returns:
      Tensor of shape [1, num_production_rules]. An element is `on_value`
      if its corresponding production rule has `symbol` on its left-hand-side;
      the element is `off_value` otherwise.
    """
    mask_values = [on_value if lhs == symbol else off_value
                   for lhs, _ in self.production_rules]
    mask_values = tf.reshape(mask_values, [1, len(self.production_rules)])
    return mask_values

class ProbabilisticGrammar(tf.keras.Model):
  """Deep generative model over productions that follow a grammar."""

  def __init__(self, grammar, latent_size, num_units):
    """Constructs a probabilistic grammar.

    Args:
      grammar: An object representing a grammar. It has members
        `nonterminal_symbols`, `alphabet`, `production_rules`, and
        `start_symbol`, and a method `mask` determining (in)valid
        production rules given a symbol.
      latent_size: Number of dimensions in the latent code.
      num_units: Number of units in the LSTM cell.
    """
    super(ProbabilisticGrammar, self).__init__()
    self.grammar = grammar
    self.latent_size = latent_size
    self.lstm = tf.compat.v1.nn.rnn_cell.LSTMCell(num_units)
    self.output_layer = tf.keras.layers.Dense(len(grammar.production_rules))

  def __call__(self, *args, **kwargs):
    inputs = 0.  # fixes a dummy variable so Model can be called without inputs
    return super(ProbabilisticGrammar, self).__call__(inputs, *args, **kwargs)

  def call(self, inputs):
    """Runs the model forward to generate a sequence of productions.

    Args:
      inputs: Unused.

    Returns:
      productions: Tensor of shape [1, num_productions, num_production_rules].
        Slices along the `num_productions` dimension represent one-hot vectors.
    """
    del inputs  # unused
    latent_code = ed.MultivariateNormalDiag(loc=tf.zeros(self.latent_size),
                                            sample_shape=1,
                                            name="latent_code")
    state = self.lstm.zero_state(1, dtype=tf.float32)
    t = 0
    productions = []
    stack = [self.grammar.start_symbol]
    while stack:
      symbol = stack.pop()
      net, state = self.lstm(latent_code, state)
      logits = (self.output_layer(net) +
                self.grammar.mask(symbol, on_value=0., off_value=-1e9))
      production = ed.OneHotCategorical(logits=logits,
                                        name="production_" + str(t))
      _, rhs = self.grammar.production_rules[tf.argmax(
          input=production, axis=-1)]
      for symbol in rhs:
        if symbol in self.grammar.nonterminal_symbols:
          stack.append(symbol)
      productions.append(production)
      t += 1
    return tf.stack(productions, axis=1)

class ProbabilisticGrammarVariational(tf.keras.Model):
  """Amortized variational posterior for a probabilistic grammar."""

  def __init__(self, latent_size):
    """Constructs a variational posterior for a probabilistic grammar.

    Args:
      latent_size: Number of dimensions in the latent code.
    """
    super(ProbabilisticGrammarVariational, self).__init__()
    self.latent_size = latent_size
    self.encoder_net = tf.keras.Sequential([
        tf.keras.layers.Conv1D(64, 3, padding="SAME"),
        tf.keras.layers.BatchNormalization(),
        tf.keras.layers.Activation(tf.nn.elu),
        tf.keras.layers.Conv1D(128, 3, padding="SAME"),
        tf.keras.layers.BatchNormalization(),
        tf.keras.layers.Activation(tf.nn.elu),
        tf.keras.layers.Dropout(0.1),
        tf.keras.layers.GlobalAveragePooling1D(),
        tf.keras.layers.Dense(latent_size * 2, activation=None),
    ])

  def call(self, inputs):
    """Runs the model forward to return a stochastic encoding.

    Args:
      inputs: Tensor of shape [1, num_productions, num_production_rules]. It is
        a sequence of productions of length `num_productions`. Each production
        is a one-hot vector of length `num_production_rules`: it determines
        which production rule the production corresponds to.

    Returns:
      latent_code_posterior: A random variable capturing a sample from the
        variational distribution, of shape [1, self.latent_size].
    """
    net = self.encoder_net(tf.cast(inputs, tf.float32))
    return ed.MultivariateNormalDiag(
        loc=net[..., :self.latent_size],
        scale_diag=tf.nn.softplus(net[..., self.latent_size:]),
        name="latent_code_posterior")

def make_value_setter(**model_kwargs):
  """Creates a value-setting interceptor.

  Args:
    **model_kwargs: dict of str to Tensor. Keys are the names of random variable
      in the model to which this interceptor is being applied. Values are
      Tensors to set their value to.

  Returns:
    set_values: Function which sets the value of intercepted ops.
  """
  def set_values(f, *args, **kwargs):
    """Sets random variable values to its aligned value."""
    name = kwargs.get("name")
    if name in model_kwargs:
      kwargs["value"] = model_kwargs[name]
    return ed.interceptable(f)(*args, **kwargs)
  return set_values

def main(argv):
  del argv  # unused
  if tf.io.gfile.exists(FLAGS.model_dir):
    tf.compat.v1.logging.warning(
        "Warning: deleting old log directory at {}".format(FLAGS.model_dir))
    tf.io.gfile.rmtree(FLAGS.model_dir)
  tf.io.gfile.makedirs(FLAGS.model_dir)
  tf.compat.v1.enable_eager_execution()

  grammar = SmilesGrammar()
  synthetic_data_distribution = ProbabilisticGrammar(
      grammar=grammar, latent_size=FLAGS.latent_size, num_units=FLAGS.num_units)

  print("Random examples from synthetic data distribution:")
  for _ in range(5):
    productions = synthetic_data_distribution()
    string = grammar.convert_to_string(productions)
    print(string)

  probabilistic_grammar = ProbabilisticGrammar(
      grammar=grammar, latent_size=FLAGS.latent_size, num_units=FLAGS.num_units)
  probabilistic_grammar_variational = ProbabilisticGrammarVariational(
      latent_size=FLAGS.latent_size)

  checkpoint = tf.train.Checkpoint(
      synthetic_data_distribution=synthetic_data_distribution,
      probabilistic_grammar=probabilistic_grammar,
      probabilistic_grammar_variational=probabilistic_grammar_variational)
  global_step = tf.compat.v1.train.get_or_create_global_step()
  optimizer = tf.compat.v1.train.AdamOptimizer(FLAGS.learning_rate)
  writer = tf.compat.v2.summary.create_file_writer(FLAGS.model_dir)
  writer.set_as_default()

  start_time = time.time()
  for step in range(FLAGS.max_steps):
    productions = synthetic_data_distribution()
    with tf.GradientTape() as tape:
      # Sample from amortized variational distribution and record its trace.
      with ed.tape() as variational_tape:
        _ = probabilistic_grammar_variational(productions)

      # Set model trace to take on the data's values and the sample from the
      # variational distribution.
      values = {"latent_code": variational_tape["latent_code_posterior"]}
      values.update({"production_" + str(t): production for t, production
                     in enumerate(tf.unstack(productions, axis=1))})
      with ed.tape() as model_tape:
        with ed.interception(make_value_setter(**values)):
          _ = probabilistic_grammar()

      # Compute the ELBO given the variational sample, averaged over the batch
      # size and the number of time steps (number of productions). Although the
      # ELBO per data point sums over time steps, we average in order to have a
      # value that remains on the same scale across batches.
      log_likelihood = 0.
      for name, rv in six.iteritems(model_tape):
        if name.startswith("production"):
          log_likelihood += rv.distribution.log_prob(rv.value)

      kl = tfp.distributions.kl_divergence(
          variational_tape["latent_code_posterior"].distribution,
          model_tape["latent_code"].distribution)

      timesteps = tf.cast(productions.shape[1], dtype=tf.float32)
      elbo = tf.reduce_mean(input_tensor=log_likelihood - kl) / timesteps
      loss = -elbo
      with tf.compat.v2.summary.record_if(
          lambda: tf.math.equal(0, global_step % 500)):
        tf.compat.v2.summary.scalar(
            "log_likelihood",
            tf.reduce_mean(input_tensor=log_likelihood) / timesteps,
            step=global_step)
        tf.compat.v2.summary.scalar(
            "kl", tf.reduce_mean(input_tensor=kl) / timesteps, step=global_step)
        tf.compat.v2.summary.scalar("elbo", elbo, step=global_step)

    variables = (probabilistic_grammar.variables
                 + probabilistic_grammar_variational.variables)
    grads = tape.gradient(loss, variables)
    grads_and_vars = zip(grads, variables)
    optimizer.apply_gradients(grads_and_vars, global_step)

    if step % 500 == 0:
      duration = time.time() - start_time
      print("Step: {:>3d} Loss: {:.3f} ({:.3f} sec)".format(
          step, loss, duration))
      checkpoint.save(file_prefix=FLAGS.model_dir)

if __name__ == "__main__":
  tf.compat.v1.app.run()