# Copyright 2019 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.
# ============================================================================
"""Tests for convolution_layers_v2.py."""
# TODO(emilyaf): Test Keras compatibility (other layers too).

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

import functools
# Dependency imports

import numpy as np
import tensorflow.compat.v2 as tf
import tensorflow_probability as tfp

from tensorflow_probability.python.internal import test_util


tfb = tfp.bijectors
tfd = tfp.distributions
tfn = tfp.experimental.nn


class BnnEndToEnd(object):

  def run_bnn_test(self, make_conv):
    # 1  Prepare Dataset

    train_size = 128
    batch_size = 2
    evidence_shape = [28, 28, 1]
    target_shape = [10]

    train_dataset = tf.data.Dataset.from_tensor_slices((
        tf.random.uniform([train_size] + evidence_shape,
                          maxval=1, dtype=tf.float32),
        tf.math.softmax(tf.random.normal([train_size] + target_shape)),
    ))
    train_dataset = tfn.util.tune_dataset(
        train_dataset,
        batch_size=batch_size,
        shuffle_size=int(train_size / 7))

    # 2  Specify Model

    scale = tfp.util.TransformedVariable(1., tfb.Softplus())
    bnn = tfn.Sequential([
        make_conv(
            evidence_shape[-1], 32, filter_shape=7,
            strides=2),     # [b, 14, 14, 32]
        tfn.util.flatten_rightmost(ndims=3),    # [b, 14 * 14 * 32]
        tfn.AffineVariationalReparameterization(
            14 * 14 * 32, np.prod(target_shape) - 1),            # [b, 9]
        lambda loc: tfb.SoftmaxCentered()(  # pylint: disable=g-long-lambda
            tfd.Independent(tfd.Normal(loc, scale),
                            reinterpreted_batch_ndims=1))  # [b, 10]
    ], name='bayesian_autoencoder')

    self.evaluate([v.initializer for v in bnn.trainable_variables])

    # 3  Train.
    n = tf.cast(train_size, tf.float32)
    train_iter = iter(train_dataset)
    def loss_fn():
      x, y = next(train_iter)
      nll = -tf.reduce_mean(bnn(x).log_prob(y), axis=-1)
      kl = tfn.losses.compute_extra_loss(bnn) / n
      return nll + kl, (nll, kl)
    opt = tf.optimizers.Adam()
    fit_op = tfn.util.make_fit_op(loss_fn, opt, bnn.trainable_variables)
    for _ in range(2):
      loss, (nll, kl) = fit_op()  # pylint: disable=unused-variable


@test_util.test_all_tf_execution_regimes
class ConvolutionV2Test(test_util.TestCase):

  def test_works_correctly(self):
    pass

  def test_non_bayesian_conv(self):
    x = tfn.ConvolutionV2(
        1, 5, 5, strides=2, padding='same')(tf.random.normal([3, 28, 28, 1]))
    self.assertShapeEqual(np.zeros([3, 14, 14, 5]), x)

  def test_bayesian_conv(self):
    def sample(kernel_shape,
               bias_shape,
               kernel_initializer=None,  # pylint: disable=unused-argument
               bias_initializer=None,  # pylint: disable=unused-argument
               kernel_batch_ndims=0,  # pylint: disable=unused-argument
               bias_batch_ndims=0,  # pylint: disable=unused-argument
               dtype=tf.float32,  # pylint: disable=unused-argument
               kernel_prior_scale_scale=None,  # pylint: disable=unused-argument
               kernel_scale_init_value=None,  # pylint: disable=unused-argument
               bias_prior_scale=None,  # pylint: disable=unused-argument
               bias_scale_init_value=None,  # pylint: disable=unused-argument
               kernel_prior=None):  # pylint: disable=unused-argument
      k = tfd.Normal(0., 1.).sample(kernel_shape)
      b = tfd.Normal(0., 1.).sample(bias_shape)
      return k, b
    def tile_for_batch(x, batch_shape):
      x_shape = tf.shape(x)
      x_ndims = tf.rank(x)
      x = tf.reshape(x, shape=tf.concat([
          x_shape[:1],
          tf.ones_like(batch_shape),
          x_shape[1:],
          ], axis=0))
      return tf.tile(x, multiples=tf.pad(
          batch_shape,
          paddings=[[1, x_ndims - 1]],
          constant_values=1))

    nn = tfn.ConvolutionV2(1, 5, 5, strides=2, padding='same', batch_shape=[10],
                           make_kernel_bias_fn=sample)
    x = tile_for_batch(tf.random.normal([3, 28, 28, 1]), [10])
    x = nn(x)
    self.assertShapeEqual(np.zeros([3, 10, 14, 14, 5]), x)


@test_util.test_all_tf_execution_regimes
class ConvolutionVariationalReparameterizationV2Test(
    test_util.TestCase, BnnEndToEnd):

  def test_works_correctly(self):
    if not tf.executing_eagerly():
      self.skipTest('Skipping graph mode test since we simulate user behavior.')
    make_conv = functools.partial(
        tfn.ConvolutionVariationalReparameterizationV2,
        rank=2,
        padding='same',
        kernel_initializer=tfn.initializers.he_uniform(),
        activation_fn=tf.nn.elu)
    self.run_bnn_test(make_conv)


@test_util.test_all_tf_execution_regimes
class ConvolutionVariationalFlipoutV2Test(test_util.TestCase, BnnEndToEnd):

  def test_works_correctly(self):
    if not tf.executing_eagerly():
      self.skipTest('Skipping graph mode test since we simulate user behavior.')
    make_conv = functools.partial(
        tfn.ConvolutionVariationalFlipoutV2,
        rank=2,
        padding='same',
        kernel_initializer=tfn.initializers.he_uniform(),
        activation_fn=tf.nn.elu)
    self.run_bnn_test(make_conv)


if __name__ == '__main__':
  test_util.main()
