text: # Copyright 2019, The TensorFlow 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.
  
  from absl.testing import parameterized
  import tensorflow.compat.v1 as tf
  from tensorflow.compat.v1.python.framework import ops as _ops
  from tensorflow.compat.v1.python.keras import keras_parameterized
  from tensorflow.compat.v1.python.keras import losses
  from tensorflow.compat.v1.python.keras.optimizer_v2.optimizer_v2 import OptimizerV2
  from tensorflow.compat.v1.python.keras.regularizers import L1L2
  from tensorflow_privacy.privacy.bolt_on import models
  from tensorflow_privacy.privacy.bolt_on.losses import StrongConvexMixin
  from tensorflow_privacy.privacy.bolt_on.optimizers import BoltOn
  
  
  class TestLoss(losses.Loss, StrongConvexMixin):
    """Test loss function for testing BoltOn model."""
  
    def __init__(self, reg_lambda, c_arg, radius_constant, name='test'):
      super().__init__(name=name)
      self.reg_lambda = reg_lambda
      self.C = c_arg  # pylint: disable=invalid-name
      self.radius_constant = radius_constant
  
    def radius(self):
      """Radius, R, of the hypothesis space W.
  
      W is a convex set that forms the hypothesis space.
  
      Returns:
        radius
      """
      return _ops.convert_to_tensor_v2(1, dtype=tf.float32)
  
    def gamma(self):
      """Returns strongly convex parameter, gamma."""
      return _ops.convert_to_tensor_v2(1, dtype=tf.float32)
  
    def beta(self, class_weight):  # pylint: disable=unused-argument
      """Smoothness, beta.
  
      Args:
        class_weight: the class weights as scalar or 1d tensor, where its
          dimensionality is equal to the number of outputs.
  
      Returns:
        Beta
      """
      return _ops.convert_to_tensor_v2(1, dtype=tf.float32)
  
    def lipchitz_constant(self, class_weight):  # pylint: disable=unused-argument
      """Lipchitz constant, L.
  
      Args:
        class_weight: class weights used
  
      Returns:
        L
      """
      return _ops.convert_to_tensor_v2(1, dtype=tf.float32)
  
    def call(self, y_true, y_pred):
      """Loss function that is minimized at the mean of the input points."""
      return 0.5 * tf.reduce_sum(
          tf.math.squared_difference(y_true, y_pred), axis=1)
  
    def max_class_weight(self, class_weight):
      """the maximum weighting in class weights (max value) as a scalar tensor.
  
      Args:
        class_weight: class weights used
  
      Returns:
        maximum class weighting as tensor scalar
      """
      if class_weight is None:
        return 1
      raise ValueError('')
  
    def kernel_regularizer(self):
      """Returns the kernel_regularizer to be used.
  
      Any subclass should override this method if they want a kernel_regularizer
      (if required for the loss function to be StronglyConvex.
      """
      return L1L2(l2=self.reg_lambda)
  
  
  class TestOptimizer(OptimizerV2):
    """Test optimizer used for testing BoltOn model."""
  
    def __init__(self):
      super().__init__('test')
  
    def compute_gradients(self):
      return 0
  
    def get_config(self):
      return {}
  
    def _create_slots(self, var):
      pass
  
    def _resource_apply_dense(self, grad, handle):
      return grad
  
    def _resource_apply_sparse(self, grad, handle, indices):
      return grad
  
  
  class InitTests(keras_parameterized.TestCase):
    """Tests for keras model initialization."""
  
    @parameterized.named_parameters([
        {
            'testcase_name': 'normal',
            'n_outputs': 1,
        },
        {
            'testcase_name': 'many outputs',
            'n_outputs': 100,
        },
    ])
    def test_init_params(self, n_outputs):
      """Test initialization of BoltOnModel.
  
      Args:
          n_outputs: number of output neurons
      """
      # test valid domains for each variable
      clf = models.BoltOnModel(n_outputs)
      self.assertIsInstance(clf, models.BoltOnModel)
  
    @parameterized.named_parameters([
        {
            'testcase_name': 'invalid n_outputs',
            'n_outputs': -1,
        },
    ])
    def test_bad_init_params(self, n_outputs):
      """test bad initializations of BoltOnModel that should raise errors.
  
      Args:
          n_outputs: number of output neurons
      """
      # test invalid domains for each variable, especially noise
      with self.assertRaises(ValueError):
        models.BoltOnModel(n_outputs)
  
    @parameterized.named_parameters([
        {
            'testcase_name': 'string compile',
            'n_outputs': 1,
            'loss': TestLoss(1, 1, 1),
            'optimizer': 'adam',
        },
        {
            'testcase_name': 'test compile',
            'n_outputs': 100,
            'loss': TestLoss(1, 1, 1),
            'optimizer': TestOptimizer(),
        },
    ])
    def test_compile(self, n_outputs, loss, optimizer):
      """Test compilation of BoltOnModel.
  
      Args:
        n_outputs: number of output neurons
        loss: instantiated TestLoss instance
        optimizer: instantiated TestOptimizer instance
      """
      # test compilation of valid tf.optimizer and tf.loss
      with self.cached_session():
        clf = models.BoltOnModel(n_outputs)
        clf.compile(optimizer, loss)
        self.assertEqual(clf.loss, loss)
  
    @parameterized.named_parameters([{
        'testcase_name': 'Not strong loss',
        'n_outputs': 1,
        'loss': losses.BinaryCrossentropy(),
        'optimizer': 'adam',
    }, {
        'testcase_name': 'Not valid optimizer',
        'n_outputs': 1,
        'loss': TestLoss(1, 1, 1),
        'optimizer': 'ada',
    }])
    def test_bad_compile(self, n_outputs, loss, optimizer):
      """test bad compilations of BoltOnModel that should raise errors.
  
      Args:
        n_outputs: number of output neurons
        loss: instantiated TestLoss instance
        optimizer: instantiated TestOptimizer instance
      """
      # test compilaton of invalid tf.optimizer and non instantiated loss.
      with self.cached_session():
        with self.assertRaises((ValueError, AttributeError)):
          clf = models.BoltOnModel(n_outputs)
          clf.compile(optimizer, loss)
  
  
  def _cat_dataset(n_samples, input_dim, n_classes, batch_size, generator=False):
    """Creates a categorically encoded dataset.
  
    Creates a categorically encoded dataset (y is categorical).
    returns the specified dataset either as a static array or as a generator.
    Will have evenly split samples across each output class.
    Each output class will be a different point in the input space.
  
    Args:
      n_samples: number of rows
      input_dim: input dimensionality
      n_classes: output dimensionality
      batch_size: The desired batch_size
      generator: False for array, True for generator
  
    Returns:
      X as (n_samples, input_dim), Y as (n_samples, n_outputs)
    """
    x_stack = []
    y_stack = []
    for i_class in range(n_classes):
      x_stack.append(tf.constant(1 * i_class, tf.float32, (n_samples, input_dim)))
      y_stack.append(tf.constant(i_class, tf.float32, (n_samples, n_classes)))
    x_set, y_set = tf.stack(x_stack), tf.stack(y_stack)
    if generator:
      dataset = tf.data.Dataset.from_tensor_slices((x_set, y_set))
      dataset = dataset.batch(batch_size=batch_size)
      return dataset
    return x_set, y_set
  
  
  def _do_fit(n_samples,
              input_dim,
              n_outputs,
              epsilon,
              generator,
              batch_size,
              reset_n_samples,
              optimizer,
              loss,
              distribution='laplace'):
    """Instantiate necessary components for fitting and perform a model fit.
  
    Args:
      n_samples: number of samples in dataset
      input_dim: the sample dimensionality
      n_outputs: number of output neurons
      epsilon: privacy parameter
      generator: True to create a generator, False to use an iterator
      batch_size: batch_size to use
      reset_n_samples: True to set _samples to None prior to fitting. False does
        nothing
      optimizer: instance of TestOptimizer
      loss: instance of TestLoss
      distribution: distribution to get noise from.
  
    Returns:
      BoltOnModel instsance
    """
    clf = models.BoltOnModel(n_outputs)
    clf.compile(optimizer, loss)
    if generator:
      x = _cat_dataset(
          n_samples, input_dim, n_outputs, batch_size, generator=generator)
      y = None
      # x = x.batch(batch_size)
      x = x.shuffle(n_samples // 2)
      batch_size = None
      if reset_n_samples:
        n_samples = None
      clf.fit_generator(
          x,
          n_samples=n_samples,
          noise_distribution=distribution,
          epsilon=epsilon)
    else:
      x, y = _cat_dataset(
          n_samples, input_dim, n_outputs, batch_size, generator=generator)
      if reset_n_samples:
        n_samples = None
      clf.fit(
          x,
          y,
          batch_size=batch_size,
          n_samples=n_samples,
          noise_distribution=distribution,
          epsilon=epsilon)
    return clf
  
  
  class FitTests(keras_parameterized.TestCase):
    """Test cases for keras model fitting."""
  
    # @test_util.run_all_in_graph_and_eager_modes
    @parameterized.named_parameters([
        {
            'testcase_name': 'iterator fit',
            'generator': False,
            'reset_n_samples': True,
        },
        {
            'testcase_name': 'iterator fit no samples',
            'generator': False,
            'reset_n_samples': True,
        },
        {
            'testcase_name': 'generator fit',
            'generator': True,
            'reset_n_samples': False,
        },
        {
            'testcase_name': 'with callbacks',
            'generator': True,
            'reset_n_samples': False,
        },
    ])
    def test_fit(self, generator, reset_n_samples):
      """Tests fitting of BoltOnModel.
  
      Args:
        generator: True for generator test, False for iterator test.
        reset_n_samples: True to reset the n_samples to None, False does nothing
      """
      loss = TestLoss(1, 1, 1)
      optimizer = BoltOn(TestOptimizer(), loss)
      n_classes = 2
      input_dim = 5
      epsilon = 1
      batch_size = 1
      n_samples = 10
      clf = _do_fit(
          n_samples,
          input_dim,
          n_classes,
          epsilon,
          generator,
          batch_size,
          reset_n_samples,
          optimizer,
          loss,
      )
      self.assertEqual(hasattr(clf, 'layers'), True)
  
    @parameterized.named_parameters([
        {
            'testcase_name': 'generator fit',
            'generator': True,
        },
    ])
    def test_fit_gen(self, generator):
      """Tests the fit_generator method of BoltOnModel.
  
      Args:
        generator: True to test with a generator dataset
      """
      loss = TestLoss(1, 1, 1)
      optimizer = TestOptimizer()
      n_classes = 2
      input_dim = 5
      batch_size = 1
      n_samples = 10
      clf = models.BoltOnModel(n_classes)
      clf.compile(optimizer, loss)
      x = _cat_dataset(
          n_samples, input_dim, n_classes, batch_size, generator=generator)
      x = x.batch(batch_size)
      x = x.shuffle(n_samples // 2)
      clf.fit_generator(x, n_samples=n_samples)
      self.assertEqual(hasattr(clf, 'layers'), True)
  
    @parameterized.named_parameters([
        {
            'testcase_name': 'iterator no n_samples',
            'generator': True,
            'reset_n_samples': True,
            'distribution': 'laplace'
        },
        {
            'testcase_name': 'invalid distribution',
            'generator': True,
            'reset_n_samples': True,
            'distribution': 'not_valid'
        },
    ])
    def test_bad_fit(self, generator, reset_n_samples, distribution):
      """Tests fitting with invalid parameters, which should raise an error.
  
      Args:
        generator: True to test with generator, False is iterator
        reset_n_samples: True to reset the n_samples param to None prior to
          passing it to fit
        distribution: distribution to get noise from.
      """
      with self.assertRaises(ValueError):
        loss = TestLoss(1, 1, 1)
        optimizer = TestOptimizer()
        n_classes = 2
        input_dim = 5
        epsilon = 1
        batch_size = 1
        n_samples = 10
        _do_fit(n_samples, input_dim, n_classes, epsilon, generator, batch_size,
                reset_n_samples, optimizer, loss, distribution)
  
    @parameterized.named_parameters([
        {
            'testcase_name': 'None class_weights',
            'class_weights': None,
            'class_counts': None,
            'num_classes': None,
            'result': 1
        },
        {
            'testcase_name': 'class weights array',
            'class_weights': [1, 1],
            'class_counts': [1, 1],
            'num_classes': 2,
            'result': [1, 1]
        },
        {
            'testcase_name': 'class weights balanced',
            'class_weights': 'balanced',
            'class_counts': [1, 1],
            'num_classes': 2,
            'result': [1, 1]
        },
    ])
    def test_class_calculate(self, class_weights, class_counts, num_classes,
                             result):
      """Tests the BOltonModel calculate_class_weights method.
  
      Args:
        class_weights: the class_weights to use
        class_counts: count of number of samples for each class
        num_classes: number of outputs neurons
        result: expected result
      """
      clf = models.BoltOnModel(1, 1)
      expected = clf.calculate_class_weights(class_weights, class_counts,
                                             num_classes)
  
      if hasattr(expected, 'numpy'):
        expected = expected.numpy()
      self.assertAllEqual(expected, result)
  
    @parameterized.named_parameters([
        {
            'testcase_name': 'class weight not valid str',
            'class_weights': 'not_valid',
            'class_counts': 1,
            'num_classes': 1,
            'err_msg': 'Detected string class_weights with value: not_valid'
        },
        {
            'testcase_name': 'no class counts',
            'class_weights': 'balanced',
            'class_counts': None,
            'num_classes': 1,
            'err_msg': 'Class counts must be provided if '
                       'using class_weights=balanced'
        },
        {
            'testcase_name': 'no num classes',
            'class_weights': 'balanced',
            'class_counts': [1],
            'num_classes': None,
            'err_msg': 'num_classes must be provided if '
                       'using class_weights=balanced'
        },
        {
            'testcase_name': 'class counts not array',
            'class_weights': 'balanced',
            'class_counts': 1,
            'num_classes': None,
            'err_msg': 'class counts must be a 1D array.'
        },
        {
            'testcase_name': 'class counts array, no num classes',
            'class_weights': [1],
            'class_counts': None,
            'num_classes': None,
            'err_msg': 'You must pass a value for num_classes if '
                       'creating an array of class_weights'
        },
        {
            'testcase_name': 'class counts array, improper shape',
            'class_weights': [[1], [1]],
            'class_counts': None,
            'num_classes': 2,
            'err_msg': 'Detected class_weights shape'
        },
        {
            'testcase_name': 'class counts array, wrong number classes',
            'class_weights': [1, 1, 1],
            'class_counts': None,
            'num_classes': 2,
            'err_msg': 'Detected array length:'
        },
    ])
    def test_class_errors(self, class_weights, class_counts, num_classes,
                          err_msg):
      """Tests the BOltonModel calculate_class_weights method.
  
      This test passes invalid params which should raise the expected errors.
  
      Args:
        class_weights: the class_weights to use.
        class_counts: count of number of samples for each class.
        num_classes: number of outputs neurons.
        err_msg: The expected error message.
      """
      clf = models.BoltOnModel(1, 1)
      with self.assertRaisesRegexp(ValueError, err_msg):  # pylint: disable=deprecated-method
        clf.calculate_class_weights(class_weights, class_counts, num_classes)
  
  
  if __name__ == '__main__':
    tf.test.main()
  