# 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_privacy.privacy.estimators import multi_class_head
from tensorflow_privacy.privacy.estimators import test_utils
from tensorflow_privacy.privacy.optimizers.dp_optimizer_keras import DPKerasSGDOptimizer
from tensorflow_estimator.python.estimator import estimator


class DPMultiClassHeadTest(tf.test.TestCase):
  """Tests for DP-enabled multiclass heads."""

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

    head = multi_class_head.DPMultiClassHead(3)
    features = {'feature_a': np.full((4), 1.0)}
    labels = np.array([[2], [1], [1], [0]])
    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.expand_dims(
        tf.compat.v1.losses.sparse_softmax_cross_entropy(
            labels=labels,
            logits=logits,
            reduction=tf.keras.losses.Reduction.NONE), -1)

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

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

    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_input_data(256, 3)
    feature_columns = []
    for key in train_features:
      feature_columns.append(tf.feature_column.numeric_column(key=key))

    head = multi_class_head.DPMultiClassHead(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 = 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_input_data(64, 3)
    classifier.evaluate(
        input_fn=test_utils.make_input_fn(test_features, test_labels, False),
        steps=4)

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


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