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

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

from tensorflow.contrib.bayesflow.python.ops import layers_conv_variational as prob_layers_lib
from tensorflow.contrib.bayesflow.python.ops import layers_util as prob_layers_util
from tensorflow.contrib.distributions.python.ops import independent as independent_lib
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor_shape
from tensorflow.python.ops import nn
from tensorflow.python.ops import nn_ops
from tensorflow.python.ops import random_ops
from tensorflow.python.ops.distributions import normal as normal_lib
from tensorflow.python.platform import test


class Counter(object):
  """Helper class to manage incrementing a counting `int`."""

  def __init__(self):
    self._value = -1

  @property
  def value(self):
    return self._value

  def __call__(self):
    self._value += 1
    return self._value


class MockDistribution(independent_lib.Independent):
  """Monitors DenseVariational calls to the underlying distribution."""

  def __init__(self, result_sample, result_log_prob, loc=None, scale=None):
    self.result_sample = result_sample
    self.result_log_prob = result_log_prob
    self.result_loc = loc
    self.result_scale = scale
    self.result_distribution = normal_lib.Normal(loc=0.0, scale=1.0)
    if loc is not None and scale is not None:
      self.result_distribution = normal_lib.Normal(loc=self.result_loc,
                                                   scale=self.result_scale)
    self.called_log_prob = Counter()
    self.called_sample = Counter()
    self.called_loc = Counter()
    self.called_scale = Counter()

  def log_prob(self, *args, **kwargs):
    self.called_log_prob()
    return self.result_log_prob

  def sample(self, *args, **kwargs):
    self.called_sample()
    return self.result_sample

  @property
  def distribution(self):  # for dummy check on Independent(Normal)
    return self.result_distribution

  @property
  def loc(self):
    self.called_loc()
    return self.result_loc

  @property
  def scale(self):
    self.called_scale()
    return self.result_scale


class MockKLDivergence(object):
  """Monitors layer calls to the divergence implementation."""

  def __init__(self, result):
    self.result = result
    self.args = []
    self.called = Counter()

  def __call__(self, *args, **kwargs):
    self.called()
    self.args.append(args)
    return self.result


class ConvVariational(test.TestCase):

  def _testKLPenaltyKernel(self, layer_class):
    with self.test_session():
      layer = layer_class(filters=2, kernel_size=3)
      if layer_class == prob_layers_lib.Conv1DVariational:
        inputs = random_ops.random_uniform([2, 3, 1], seed=1)
      elif layer_class == prob_layers_lib.Conv2DVariational:
        inputs = random_ops.random_uniform([2, 3, 3, 1], seed=1)
      elif layer_class == prob_layers_lib.Conv3DVariational:
        inputs = random_ops.random_uniform([2, 3, 3, 3, 1], seed=1)

      # No keys.
      losses = ops.get_collection(ops.GraphKeys.REGULARIZATION_LOSSES)
      self.assertEqual(len(losses), 0)
      self.assertListEqual(layer.losses, losses)

      _ = layer(inputs)

      # Yes keys.
      losses = ops.get_collection(ops.GraphKeys.REGULARIZATION_LOSSES)
      self.assertEqual(len(losses), 1)
      self.assertListEqual(layer.losses, losses)

  def _testKLPenaltyBoth(self, layer_class):
    def _make_normal(dtype, *args):  # pylint: disable=unused-argument
      return normal_lib.Normal(
          loc=dtype.as_numpy_dtype(0.), scale=dtype.as_numpy_dtype(1.))
    with self.test_session():
      layer = layer_class(
          filters=2,
          kernel_size=3,
          bias_posterior_fn=prob_layers_util.default_mean_field_normal_fn(),
          bias_prior_fn=_make_normal)
      if layer_class == prob_layers_lib.Conv1DVariational:
        inputs = random_ops.random_uniform([2, 3, 1], seed=1)
      elif layer_class == prob_layers_lib.Conv2DVariational:
        inputs = random_ops.random_uniform([2, 3, 3, 1], seed=1)
      elif layer_class == prob_layers_lib.Conv3DVariational:
        inputs = random_ops.random_uniform([2, 3, 3, 3, 1], seed=1)

      # No keys.
      losses = ops.get_collection(ops.GraphKeys.REGULARIZATION_LOSSES)
      self.assertEqual(len(losses), 0)
      self.assertListEqual(layer.losses, losses)

      _ = layer(inputs)

      # Yes keys.
      losses = ops.get_collection(ops.GraphKeys.REGULARIZATION_LOSSES)
      self.assertEqual(len(losses), 2)
      self.assertListEqual(layer.losses, losses)

  def _testConvVariational(self, layer_class):
    batch_size, depth, height, width, channels, filters = 2, 4, 4, 4, 3, 5
    with self.test_session() as sess:
      seed = Counter()
      if layer_class == prob_layers_lib.Conv1DVariational:
        inputs = random_ops.random_uniform(
            [batch_size, width, channels], seed=seed())
        kernel_size = (2,)
      elif layer_class == prob_layers_lib.Conv2DVariational:
        inputs = random_ops.random_uniform(
            [batch_size, height, width, channels], seed=seed())
        kernel_size = (2, 2)
      elif layer_class == prob_layers_lib.Conv3DVariational:
        inputs = random_ops.random_uniform(
            [batch_size, depth, height, width, channels], seed=seed())
        kernel_size = (2, 2, 2)

      kernel_shape = kernel_size + (channels, filters)
      kernel_posterior = MockDistribution(
          result_log_prob=random_ops.random_uniform(kernel_shape, seed=seed()),
          result_sample=random_ops.random_uniform(kernel_shape, seed=seed()))
      kernel_prior = MockDistribution(
          result_log_prob=random_ops.random_uniform(kernel_shape, seed=seed()),
          result_sample=random_ops.random_uniform(kernel_shape, seed=seed()))
      kernel_divergence = MockKLDivergence(
          result=random_ops.random_uniform(kernel_shape, seed=seed()))

      bias_size = (filters,)
      bias_posterior = MockDistribution(
          result_log_prob=random_ops.random_uniform(bias_size, seed=seed()),
          result_sample=random_ops.random_uniform(bias_size, seed=seed()))
      bias_prior = MockDistribution(
          result_log_prob=random_ops.random_uniform(bias_size, seed=seed()),