# Copyright 2018, 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.
"""Unit testing for loss.py"""

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

import tensorflow as tf
from tensorflow.python.platform import test
from tensorflow.python.keras import keras_parameterized
from tensorflow.python.keras.optimizer_v2 import adam
from tensorflow.python.keras.optimizer_v2 import adagrad
from tensorflow.python.keras.optimizer_v2 import gradient_descent
from tensorflow.python.keras import losses
from tensorflow.python.framework import test_util
from privacy.bolton import model
from privacy.bolton.loss import StrongConvexBinaryCrossentropy
from privacy.bolton.loss import StrongConvexHuber
from privacy.bolton.loss import StrongConvexMixin
from absl.testing import parameterized
from absl.testing import absltest
from tensorflow.python.keras.regularizers import L1L2


class StrongConvexTests(keras_parameterized.TestCase):
  @parameterized.named_parameters([
      {'testcase_name': 'beta not implemented',
       'fn': 'beta',
       'args': [1]},
      {'testcase_name': 'gamma not implemented',
       'fn': 'gamma',
       'args': []},
      {'testcase_name': 'lipchitz not implemented',
       'fn': 'lipchitz_constant',
       'args': [1]},
      {'testcase_name': 'radius not implemented',
       'fn': 'radius',
       'args': []},
  ])
  def test_not_implemented(self, fn, args):
    with self.assertRaises(NotImplementedError):
      loss = StrongConvexMixin()
      getattr(loss, fn, None)(*args)

  @parameterized.named_parameters([
      {'testcase_name': 'radius not implemented',
       'fn': 'kernel_regularizer',
       'args': []},
  ])
  def test_return_none(self, fn, args):
    loss = StrongConvexMixin()
    ret = getattr(loss, fn, None)(*args)
    self.assertEqual(ret, None)


class BinaryCrossesntropyTests(keras_parameterized.TestCase):
  """tests for BinaryCrossesntropy StrongConvex loss"""

  @parameterized.named_parameters([
      {'testcase_name': 'normal',
       'reg_lambda': 1,
       'c': 1,
       'radius_constant': 1
       },
  ])
  def test_init_params(self, reg_lambda, c, radius_constant):
    # test valid domains for each variable
    loss = StrongConvexBinaryCrossentropy(reg_lambda, c, radius_constant)
    self.assertIsInstance(loss, StrongConvexBinaryCrossentropy)

  @parameterized.named_parameters([
      {'testcase_name': 'negative c',
       'reg_lambda': 1,
       'c': -1,
       'radius_constant': 1
       },
      {'testcase_name': 'negative radius',
       'reg_lambda': 1,
       'c': 1,
       'radius_constant': -1
       },
      {'testcase_name': 'negative lambda',
       'reg_lambda': -1,
       'c': 1,
       'radius_constant': 1
       },
  ])
  def test_bad_init_params(self, reg_lambda, c, radius_constant):
    # test valid domains for each variable
    with self.assertRaises(ValueError):
      loss = StrongConvexBinaryCrossentropy(reg_lambda, c, radius_constant)

  @test_util.run_all_in_graph_and_eager_modes
  @parameterized.named_parameters([
      # [] for compatibility with tensorflow loss calculation
      {'testcase_name': 'both positive',
        'logits': [10000],
        'y_true': [1],
        'result': 0,
      },
      {'testcase_name': 'positive gradient negative logits',
       'logits': [-10000],
       'y_true': [1],
       'result': 10000,
       },
      {'testcase_name': 'positivee gradient positive logits',
       'logits': [10000],
       'y_true': [0],
       'result': 10000,
       },
      {'testcase_name': 'both negative',
       'logits': [-10000],
       'y_true': [0],
       'result': 0
       },
  ])
  def test_calculation(self, logits, y_true, result):
    logits = tf.Variable(logits, False, dtype=tf.float32)
    y_true = tf.Variable(y_true, False, dtype=tf.float32)
    loss = StrongConvexBinaryCrossentropy(0.00001, 1, 1)
    loss = loss(y_true, logits)
    self.assertEqual(loss.numpy(), result)

  @parameterized.named_parameters([
      {'testcase_name': 'beta',
       'init_args': [1, 1, 1],
       'fn': 'beta',
       'args': [1],
       'result': tf.constant(2, dtype=tf.float32)
       },
      {'testcase_name': 'gamma',
       'fn': 'gamma',
       'init_args': [1, 1, 1],
       'args': [],
       'result': tf.constant(1, dtype=tf.float32),
       },
      {'testcase_name': 'lipchitz constant',
       'fn': 'lipchitz_constant',
       'init_args': [1, 1, 1],
       'args': [1],
       'result': tf.constant(2, dtype=tf.float32),
       },
      {'testcase_name': 'kernel regularizer',
       'fn': 'kernel_regularizer',
       'init_args': [1, 1, 1],
       'args': [],
       'result': L1L2(l2=1),
       },
  ])
  def test_fns(self, init_args, fn, args, result):
    loss = StrongConvexBinaryCrossentropy(*init_args)
    expected = getattr(loss, fn, lambda: 'fn not found')(*args)
    if hasattr(expected, 'numpy') and hasattr(result, 'numpy'):  # both tensor
      expected = expected.numpy()
      result = result.numpy()
    if hasattr(expected, 'l2') and hasattr(result, 'l2'):  # both l2 regularizer
      expected = expected.l2
      result = result.l2
    self.assertEqual(expected, result)


class HuberTests(keras_parameterized.TestCase):
  """tests for BinaryCrossesntropy StrongConvex loss"""

  @parameterized.named_parameters([
      {'testcase_name': 'normal',
       'reg_lambda': 1,
       'c': 1,
       'radius_constant': 1,
       'delta': 1,
       },
  ])
  def test_init_params(self, reg_lambda, c, radius_constant, delta):
    # test valid domains for each variable
    loss = StrongConvexHuber(reg_lambda, c, radius_constant, delta)
    self.assertIsInstance(loss, StrongConvexHuber)

  @parameterized.named_parameters([
      {'testcase_name': 'negative c',
       'reg_lambda': 1,
       'c': -1,
       'radius_constant': 1,
       'delta': 1
       },
      {'testcase_name': 'negative radius',
       'reg_lambda': 1,
       'c': 1,
       'radius_constant': -1,
       'delta': 1
       },
      {'testcase_name': 'negative lambda',
       'reg_lambda': -1,
       'c': 1,
       'radius_constant': 1,
       'delta': 1
       },
      {'testcase_name': 'negative delta',
       'reg_lambda': -1,
       'c': 1,
       'radius_constant': 1,
       'delta': -1
       },
  ])
  def test_bad_init_params(self, reg_lambda, c, radius_constant, delta):
    # test valid domains for each variable
    with self.assertRaises(ValueError):
      loss = StrongConvexHuber(reg_lambda, c, radius_constant, delta)

  # test the bounds and test varied delta's
  @test_util.run_all_in_graph_and_eager_modes
  @parameterized.named_parameters([
      {'testcase_name': 'delta=1,y_true=1 z>1+h decision boundary',
        'logits': 2.1,
        'y_true': 1,
        'delta': 1,
        'result': 0,
      },
      {'testcase_name': 'delta=1,y_true=1 z<1+h decision boundary',
       'logits': 1.9,
       'y_true': 1,
       'delta': 1,
       'result': 0.01*0.25,
       },
      {'testcase_name': 'delta=1,y_true=1 1-z< h decision boundary',
       'logits': 0.1,
       'y_true': 1,
       'delta': 1,
       'result': 1.9**2 * 0.25,
       },
      {'testcase_name': 'delta=1,y_true=1 z < 1-h decision boundary',
       'logits': -0.1,
       'y_true': 1,
       'delta': 1,
       'result': 1.1,
       },
      {'testcase_name': 'delta=2,y_true=1 z>1+h decision boundary',
       'logits': 3.1,
       'y_true': 1,
       'delta': 2,
       'result': 0,
       },
      {'testcase_name': 'delta=2,y_true=1 z<1+h decision boundary',
       'logits': 2.9,
       'y_true': 1,
       'delta': 2,
       'result': 0.01*0.125,
       },
      {'testcase_name': 'delta=2,y_true=1 1-z < h decision boundary',
       'logits': 1.1,
       'y_true': 1,
       'delta': 2,
       'result': 1.9**2 * 0.125,
       },
      {'testcase_name': 'delta=2,y_true=1 z < 1-h decision boundary',
       'logits': -1.1,
       'y_true': 1,
       'delta': 2,
       'result': 2.1,
       },
      {'testcase_name': 'delta=1,y_true=-1 z>1+h decision boundary',
       'logits': -2.1,
       'y_true': -1,
       'delta': 1,
       'result': 0,
       },
  ])
  def test_calculation(self, logits, y_true, delta, result):
    logits = tf.Variable(logits, False, dtype=tf.float32)
    y_true = tf.Variable(y_true, False, dtype=tf.float32)
    loss = StrongConvexHuber(0.00001, 1, 1, delta)
    loss = loss(y_true, logits)
    self.assertAllClose(loss.numpy(), result)

  @parameterized.named_parameters([
      {'testcase_name': 'beta',
       'init_args': [1, 1, 1, 1],
       'fn': 'beta',
       'args': [1],
       'result': tf.Variable(1.5, dtype=tf.float32)
       },
      {'testcase_name': 'gamma',
       'fn': 'gamma',
       'init_args': [1, 1, 1, 1],
       'args': [],
       'result': tf.Variable(1, dtype=tf.float32),
       },
      {'testcase_name': 'lipchitz constant',
       'fn': 'lipchitz_constant',
       'init_args': [1, 1, 1, 1],
       'args': [1],
       'result': tf.Variable(2, dtype=tf.float32),
       },
      {'testcase_name': 'kernel regularizer',
       'fn': 'kernel_regularizer',
       'init_args': [1, 1, 1, 1],
       'args': [],
       'result': L1L2(l2=1),
       },
  ])
  def test_fns(self, init_args, fn, args, result):
    loss = StrongConvexHuber(*init_args)
    expected = getattr(loss, fn, lambda: 'fn not found')(*args)
    if hasattr(expected, 'numpy') and hasattr(result, 'numpy'):  # both tensor
      expected = expected.numpy()
      result = result.numpy()
    if hasattr(expected, 'l2') and hasattr(result, 'l2'):  # both l2 regularizer
      expected = expected.l2
      result = result.l2
    self.assertEqual(expected, result)


if __name__ == '__main__':
  tf.test.main()