# Copyright 2020, 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.

import numpy as np
import tensorflow as tf
from tensorflow import estimator as tf_estimator
from tensorflow_privacy.privacy.estimators import multi_label_head
from tensorflow_privacy.privacy.estimators import test_utils
from tensorflow_privacy.privacy.optimizers.dp_optimizer_keras import DPKerasSGDOptimizer


class DPMultiLabelHeadTest(tf.test.TestCase):
  """Tests for DP-enabled multilabel heads."""

  def testLoss(self):
    """Tests loss() returns per-example losses."""

    head = multi_label_head.DPMultiLabelHead(3)
    features = {'feature_a': np.full((4), 1.0)}
    labels = np.array([[0, 1, 1], [1, 1, 0], [0, 1, 0], [1, 1, 1]])
    logits = np.array([[2.0, 1.5, 4.1], [2.0, 1.5, 4.1], [2.0, 1.5, 4.1],
                       [2.0, 1.5, 4.1]])

    actual_loss = head.loss(labels, logits, features)
    expected_loss = tf.reduce_mean(
        tf.compat.v1.losses.sigmoid_cross_entropy(
            multi_class_labels=labels,
            logits=logits,
            reduction=tf.keras.losses.Reduction.NONE),
        axis=-1,
        keepdims=True)

    if tf.executing_eagerly():
      self.assertEqual(actual_loss.shape, [4, 1])
      self.assertAllClose(actual_loss, expected_loss)
      return

    self.assertEqual(actual_loss.shape, [4, 1])
    self.assertAllClose(expected_loss, self.evaluate(actual_loss))

  def testCreateTPUEstimatorSpec(self):
    """Tests that an Estimator built with this head works."""

    train_features, train_labels = test_utils.make_multilabel_input_data(256)
    feature_columns = []
    for key in train_features:
      feature_columns.append(tf.feature_column.numeric_column(key=key))

    head = multi_label_head.DPMultiLabelHead(3)
    optimizer = DPKerasSGDOptimizer(
        learning_rate=0.5,
        l2_norm_clip=1.0,
        noise_multiplier=0.0,
        num_microbatches=2)
    model_fn = test_utils.make_model_fn(head, optimizer, feature_columns)
    classifier = tf_estimator.Estimator(model_fn=model_fn)

    classifier.train(
        input_fn=test_utils.make_input_fn(train_features, train_labels, True),
        steps=4)

    test_features, test_labels = test_utils.make_multilabel_input_data(64)
    classifier.evaluate(
        input_fn=test_utils.make_input_fn(test_features, test_labels, False),
        steps=4)

    predict_features, predict_labels = test_utils.make_multilabel_input_data(64)
    classifier.predict(
        input_fn=test_utils.make_input_fn(predict_features, predict_labels,
                                          False))


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