# 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.
# ============================================================================
"""DenseVariational layer."""

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

import tensorflow.compat.v2 as tf

from tensorflow_probability.python.distributions import kullback_leibler


class DenseVariational(tf.keras.layers.Layer):
  """Dense layer with random `kernel` and `bias`.

  This layer uses variational inference to fit a "surrogate" posterior to the
  distribution over both the `kernel` matrix and the `bias` terms which are
  otherwise used in a manner similar to `tf.keras.layers.Dense`.

  This layer fits the "weights posterior" according to the following generative
  process:

  ```none
  [K, b] ~ Prior()
  M = matmul(X, K) + b
  Y ~ Likelihood(M)
  ```

  """

  def __init__(self,
               units,
               make_posterior_fn,
               make_prior_fn,
               kl_weight=None,
               kl_use_exact=False,
               activation=None,
               use_bias=True,
               activity_regularizer=None,
               **kwargs):
    """Creates the `DenseVariational` layer.

    Args:
      units: Positive integer, dimensionality of the output space.
      make_posterior_fn: Python callable taking `tf.size(kernel)`,
        `tf.size(bias)`, `dtype` and returns another callable which takes an
        input and produces a `tfd.Distribution` instance.
      make_prior_fn: Python callable taking `tf.size(kernel)`, `tf.size(bias)`,
        `dtype` and returns another callable which takes an input and produces a
        `tfd.Distribution` instance.
      kl_weight: Amount by which to scale the KL divergence loss between prior
        and posterior.
      kl_use_exact: Python `bool` indicating that the analytical KL divergence
        should be used rather than a Monte Carlo approximation.
      activation: Activation function to use.
        If you don't specify anything, no activation is applied
        (ie. "linear" activation: `a(x) = x`).
      use_bias: Boolean, whether the layer uses a bias vector.
      activity_regularizer: Regularizer function applied to
        the output of the layer (its "activation")..
      **kwargs: Extra arguments forwarded to `tf.keras.layers.Layer`.
    """
    if 'input_shape' not in kwargs and 'input_dim' in kwargs:
      kwargs['input_shape'] = (kwargs.pop('input_dim'),)
    super(DenseVariational, self).__init__(
        activity_regularizer=tf.keras.regularizers.get(activity_regularizer),
        **kwargs)
    self.units = int(units)

    self._make_posterior_fn = make_posterior_fn
    self._make_prior_fn = make_prior_fn
    self._kl_divergence_fn = _make_kl_divergence_penalty(
        kl_use_exact, weight=kl_weight)

    self.activation = tf.keras.activations.get(activation)
    self.use_bias = use_bias
    self.supports_masking = False
    self.input_spec = tf.keras.layers.InputSpec(min_ndim=2)

  def build(self, input_shape):
    dtype = tf.as_dtype(self.dtype or tf.keras.backend.floatx())
    if not (dtype.is_floating or dtype.is_complex):
      raise TypeError('Unable to build `Dense` layer with non-floating point '
                      'dtype %s' % (dtype,))
    input_shape = tf.TensorShape(input_shape)
    last_dim = tf.compat.dimension_value(input_shape[-1])
    if last_dim is None:
      raise ValueError('The last dimension of the inputs to `DenseVariational` '
                       'should be defined. Found `None`.')
    self.input_spec = tf.keras.layers.InputSpec(
        min_ndim=2, axes={-1: last_dim})

    self._posterior = self._make_posterior_fn(
        last_dim * self.units,
        self.units if self.use_bias else 0,
        dtype)
    self._prior = self._make_prior_fn(
        last_dim * self.units,
        self.units if self.use_bias else 0,
        dtype)

    self.built = True

  def call(self, inputs):
    dtype = tf.as_dtype(self.dtype or tf.keras.backend.floatx())
    inputs = tf.cast(inputs, dtype, name='inputs')

    q = self._posterior(inputs)
    r = self._prior(inputs)
    self.add_loss(self._kl_divergence_fn(q, r))

    w = tf.convert_to_tensor(value=q)
    prev_units = self.input_spec.axes[-1]
    if self.use_bias:
      split_sizes = [prev_units * self.units, self.units]
      kernel, bias = tf.split(w, split_sizes, axis=-1)
    else:
      kernel, bias = w, None

    kernel = tf.reshape(kernel, shape=tf.concat([
        tf.shape(kernel)[:-1],
        [prev_units, self.units],
    ], axis=0))
    outputs = tf.matmul(inputs, kernel)

    if self.use_bias:
      outputs = tf.nn.bias_add(outputs, bias)

    if self.activation is not None:
      outputs = self.activation(outputs)  # pylint: disable=not-callable

    return outputs


def _make_kl_divergence_penalty(
    use_exact_kl=False,
    test_points_reduce_axis=(),  # `None` == "all"; () == "none".
    test_points_fn=tf.convert_to_tensor,
    weight=None):
  """Creates a callable computing `KL[a,b]` from `a`, a `tfd.Distribution`."""

  if use_exact_kl:
    kl_divergence_fn = kullback_leibler.kl_divergence
  else:
    def kl_divergence_fn(distribution_a, distribution_b):
      z = test_points_fn(distribution_a)
      return tf.reduce_mean(
          distribution_a.log_prob(z) - distribution_b.log_prob(z),
          axis=test_points_reduce_axis)

  # Closure over: kl_divergence_fn, weight.
  def _fn(distribution_a, distribution_b):
    """Closure that computes KLDiv as a function of `a` as in `KL[a, b]`."""
    with tf.name_scope('kldivergence_loss'):
      kl = kl_divergence_fn(distribution_a, distribution_b)
      if weight is not None:
        kl = tf.cast(weight, dtype=kl.dtype) * kl
      # Losses appended with the model.add_loss and are expected to be a single
      # scalar, unlike model.loss, which is expected to be the loss per sample.
      # Therefore, we reduce over all dimensions, regardless of the shape.
      # We take the sum because (apparently) Keras will add this to the *post*
      # `reduce_sum` (total) loss.
      # TODO(b/126259176): Add end-to-end Keras/TFP test to ensure the API's
      # align, particularly wrt how losses are aggregated (across batch
      # members).
      return tf.reduce_sum(kl, name='batch_total_kl_divergence')

  return _fn
