# Copyright 2018 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 binary_class_head.py."""

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

import numpy as np
import six
import tensorflow as tf
from tensorflow.python.framework import test_util
from tensorflow_estimator.python.estimator.canned import dnn
from tensorflow_estimator.python.estimator.canned import dnn_testing_utils
from tensorflow_estimator.python.estimator.canned import metric_keys
from tensorflow_estimator.python.estimator.canned import prediction_keys
from tensorflow_estimator.python.estimator.head import binary_class_head as head_lib
from tensorflow_estimator.python.estimator.head import head_utils as test_lib
from tensorflow_estimator.python.estimator.mode_keys import ModeKeys


@test_util.run_all_in_graph_and_eager_modes
class BinaryClassHeadTest(tf.test.TestCase):

  def test_threshold_too_small(self):
    with self.assertRaisesRegexp(ValueError, r'thresholds not in \(0, 1\)'):
      head_lib.BinaryClassHead(thresholds=(0., 0.5))

  def test_threshold_too_large(self):
    with self.assertRaisesRegexp(ValueError, r'thresholds not in \(0, 1\)'):
      head_lib.BinaryClassHead(thresholds=(0.5, 1.))

  def test_invalid_loss_reduction(self):
    with self.assertRaisesRegexp(
        ValueError, r'Invalid loss_reduction: invalid_loss_reduction'):
      head_lib.BinaryClassHead(loss_reduction='invalid_loss_reduction')
    with self.assertRaisesRegexp(ValueError, r'Invalid loss_reduction: none'):
      head_lib.BinaryClassHead(loss_reduction=tf.losses.Reduction.NONE)

  def test_loss_fn_arg_labels_missing(self):

    def _loss_fn(logits):
      del logits  # Unused

    with self.assertRaisesRegexp(
        ValueError, r'loss_fn must contain argument: labels\. '
        r'Given arguments: \(\'logits\',\)'):
      head_lib.BinaryClassHead(loss_fn=_loss_fn)

  def test_loss_fn_arg_logits_missing(self):

    def _loss_fn(labels):
      del labels  # unused

    with self.assertRaisesRegexp(
        ValueError, r'loss_fn must contain argument: logits\. '
        r'Given arguments: \(\'labels\',\)'):
      head_lib.BinaryClassHead(loss_fn=_loss_fn)

  def test_loss_fn_arg_features_ok(self):

    def _loss_fn(labels, logits, features):
      del labels, logits, features  # Unused
      head_lib.BinaryClassHead(loss_fn=_loss_fn)

  def test_loss_fn_arg_invalid(self):

    def _loss_fn(labels, logits, name=None):
      del labels, logits, name  # Unused

    with self.assertRaisesRegexp(ValueError,
                                 r'loss_fn has unexpected args: \[\'name\'\]'):
      head_lib.BinaryClassHead(loss_fn=_loss_fn)

  def test_invalid_logits_shape(self):
    head = head_lib.BinaryClassHead()
    self.assertEqual(1, head.logits_dimension)

    # Logits should be shape (batch_size, 1).
    logits_2x2 = np.array(((45., 44.), (41., 42.),))

    pred_key = prediction_keys.PredictionKeys.PROBABILITIES
    # Static shape.
    with self.assertRaisesRegexp(ValueError, 'logits shape'):
      preds = head.predictions(logits_2x2, [pred_key])
      self.evaluate(preds[pred_key])
    if tf.executing_eagerly():
      return

    # Dynamic shape only works in Graph mode.
    logits_placeholder = tf.compat.v1.placeholder(dtype=tf.dtypes.float32)
    spec = head.create_estimator_spec(
        features={'x': np.array(((42.,),))},
        mode=ModeKeys.PREDICT,
        logits=logits_placeholder,
        trainable_variables=[tf.Variable([1.0, 2.0], dtype=tf.dtypes.float32)])
    with self.cached_session():
      with self.assertRaisesRegexp(tf.errors.OpError, 'logits shape'):
        spec.predictions[pred_key].eval({logits_placeholder: logits_2x2})

  def test_invalid_labels_shape(self):
    head = head_lib.BinaryClassHead()
    self.assertEqual(1, head.logits_dimension)

    # Labels and logits should be shape (batch_size, 1).
    labels_2x2 = np.array(((45., 44.), (41., 42.),))
    logits_2x1 = np.array(((45.,), (41.,),))
    features = {'x': np.array(((42.,),))}

    # Static shape.
    with self.assertRaisesRegexp(ValueError, 'Mismatched label shape'):
      training_loss = head.loss(
          logits=logits_2x1,
          labels=labels_2x2,
          features=features,
          mode=ModeKeys.EVAL)
      self.evaluate(training_loss)
    if tf.executing_eagerly():
      return

    # Dynamic shape only works in Graph mode.
    labels_placeholder = tf.compat.v1.placeholder(dtype=tf.dtypes.float32)
    logits_placeholder = tf.compat.v1.placeholder(dtype=tf.dtypes.float32)
    training_loss = head.loss(
        logits=logits_placeholder,
        labels=labels_placeholder,
        features=features,
        mode=ModeKeys.EVAL)
    with self.cached_session():
      with self.assertRaisesRegexp(
          tf.errors.InvalidArgumentError,
          r'\[expected_labels_shape: \] \[2 1\] \[labels_shape: \] \[2 2\]'):
        training_loss.eval({
            logits_placeholder: logits_2x1,
            labels_placeholder: labels_2x2
        })

  def test_incompatible_labels_shape(self):
    head = head_lib.BinaryClassHead()
    self.assertEqual(1, head.logits_dimension)

    # Both logits and labels should be shape (batch_size, 1).
    values_2x1 = np.array(((0.,), (1.,),))
    values_3x1 = np.array(((0.,), (1.,), (0.,),))
    features = {'x': values_2x1}

    # Static shape for eager mode.
    if tf.executing_eagerly():
      with self.assertRaisesRegexp(ValueError, 'labels shape'):
        head.loss(
            logits=values_2x1,
            labels=values_3x1,
            features=features,
            mode=ModeKeys.EVAL)
      with self.assertRaisesRegexp(ValueError, 'labels shape'):
        head.loss(
            logits=values_3x1,
            labels=values_2x1,
            features=features,
            mode=ModeKeys.EVAL)
      return

    # Static shape for Graph mode.
    with self.assertRaisesRegexp(ValueError,
                                 'logits and labels must have the same shape'):
      head.loss(
          logits=values_2x1,
          labels=values_3x1,
          features=features,
          mode=ModeKeys.EVAL)
    with self.assertRaisesRegexp(ValueError,
                                 'logits and labels must have the same shape'):
      head.loss(
          logits=values_3x1,
          labels=values_2x1,
          features=features,
          mode=ModeKeys.EVAL)
    # Dynamic shape only works in Graph mode.
    labels_placeholder = tf.compat.v1.placeholder(dtype=tf.dtypes.float32)
    logits_placeholder = tf.compat.v1.placeholder(dtype=tf.dtypes.float32)
    training_loss = head.loss(
        logits=logits_placeholder,
        labels=labels_placeholder,
        features=features,
        mode=ModeKeys.EVAL)
    with self.cached_session():
      with self.assertRaisesRegexp(
          tf.errors.InvalidArgumentError,
          r'\[expected_labels_shape: \] \[3 1\] \[labels_shape: \] \[2 1\]'):
        training_loss.eval({
            labels_placeholder: values_2x1,
            logits_placeholder: values_3x1
        })
    with self.cached_session():
      with self.assertRaisesRegexp(
          tf.errors.InvalidArgumentError,
          r'\[expected_labels_shape: \] \[2 1\] \[labels_shape: \] \[3 1\]'):
        training_loss.eval({
            labels_placeholder: values_3x1,
            logits_placeholder: values_2x1
        })

  def test_predict(self):
    head = head_lib.BinaryClassHead()
    self.assertEqual(1, head.logits_dimension)

    logits = [[0.3], [-0.4]]
    expected_logistics = [[0.574443], [0.401312]]
    expected_probabilities = [[0.425557, 0.574443], [0.598688, 0.401312]]
    expected_class_ids = [[1], [0]]
    expected_classes = [[b'1'], [b'0']]
    expected_all_class_ids = [[0, 1]] * 2
    expected_all_classes = [[b'0', b'1']] * 2
    expected_export_classes = [[b'0', b'1']] * 2

    keys = prediction_keys.PredictionKeys
    preds = head.predictions(logits, [
        keys.LOGITS, keys.LOGISTIC, keys.PROBABILITIES, keys.CLASS_IDS,
        keys.CLASSES, keys.ALL_CLASS_IDS, keys.ALL_CLASSES
    ])
    self.assertAllClose(logits, self.evaluate(preds[keys.LOGITS]))
    self.assertAllClose(expected_logistics, self.evaluate(preds[keys.LOGISTIC]))
    self.assertAllClose(expected_probabilities,
                        self.evaluate(preds[keys.PROBABILITIES]))
    self.assertAllClose(expected_class_ids,
                        self.evaluate(preds[keys.CLASS_IDS]))
    self.assertAllClose(expected_all_class_ids,
                        self.evaluate(preds[keys.ALL_CLASS_IDS]))
    self.assertAllEqual(expected_classes, self.evaluate(preds[keys.CLASSES]))
    self.assertAllEqual(expected_all_classes,
                        self.evaluate(preds[keys.ALL_CLASSES]))
    if tf.executing_eagerly():
      return

    # Create estimator spec.
    spec = head.create_estimator_spec(
        features={'x': np.array(((42,),), dtype=np.int32)},
        mode=ModeKeys.PREDICT,
        logits=logits,
        trainable_variables=[tf.Variable([1.0, 2.0], dtype=tf.dtypes.float32)])

    # Assert spec contains expected tensors.
    self.assertIsNone(spec.loss)
    self.assertEqual({}, spec.eval_metric_ops)
    self.assertIsNone(spec.train_op)
    self.assertItemsEqual(('classification', 'regression', 'predict',
                           test_lib._DEFAULT_SERVING_KEY),
                          spec.export_outputs.keys())
    test_lib._assert_no_hooks(self, spec)

    # Assert predictions.
    with self.cached_session() as sess:
      test_lib._initialize_variables(self, spec.scaffold)
      self.assertIsNone(spec.scaffold.summary_op)
      predictions = sess.run(spec.predictions)
      self.assertAllClose(logits,
                          predictions[prediction_keys.PredictionKeys.LOGITS])
      self.assertAllClose(expected_logistics,
                          predictions[prediction_keys.PredictionKeys.LOGISTIC])
      self.assertAllClose(
          expected_probabilities,
          predictions[prediction_keys.PredictionKeys.PROBABILITIES])
      self.assertAllClose(expected_class_ids,
                          predictions[prediction_keys.PredictionKeys.CLASS_IDS])
      self.assertAllEqual(expected_classes,
                          predictions[prediction_keys.PredictionKeys.CLASSES])
      self.assertAllClose(
          expected_all_class_ids,
          predictions[prediction_keys.PredictionKeys.ALL_CLASS_IDS])
      self.assertAllEqual(
          expected_all_classes,
          predictions[prediction_keys.PredictionKeys.ALL_CLASSES])
      self.assertAllClose(
          expected_probabilities,
          sess.run(spec.export_outputs[test_lib._DEFAULT_SERVING_KEY].scores))
      self.assertAllEqual(
          expected_export_classes,
          sess.run(spec.export_outputs[test_lib._DEFAULT_SERVING_KEY].classes))
      self.assertAllClose(expected_logistics,
                          sess.run(spec.export_outputs['regression'].value))

  def test_predict_with_vocabulary_list(self):
    head = head_lib.BinaryClassHead(label_vocabulary=['aang', 'iroh'])

    logits = [[1.], [0.]]
    expected_classes = [[b'iroh'], [b'aang']]

    pred_key = prediction_keys.PredictionKeys.CLASSES
    if tf.executing_eagerly():
      preds = head.predictions(logits, [pred_key])
      self.assertAllEqual(expected_classes, preds[pred_key])
      return
    preds = head.predictions(logits, [pred_key])
    with self.cached_session():
      test_lib._initialize_variables(self, tf.compat.v1.train.Scaffold())
      self.assertAllEqual(expected_classes, preds[pred_key].eval())

  def test_eval_create_loss(self):
    head = head_lib.BinaryClassHead()
    logits = np.array(((45,), (-41,),), dtype=np.float32)
    labels = np.array(((1,), (1,),), dtype=np.int32)
    features = {'x': np.array(((42,),), dtype=np.int32)}

    # loss = sum(cross_entropy(labels, logits)) / batch_size
    #      = sum([0, 41]) / 2 = 20.5
    expected_training_loss = 20.5
    # Create loss.
    training_loss = head.loss(
        logits=logits, labels=labels, features=features, mode=ModeKeys.EVAL)
    self.assertAllClose(
        expected_training_loss,
        self.evaluate(training_loss),
        rtol=1e-2,
        atol=1e-2)

  def test_eval_create_loss_loss_fn(self):
    """Tests head.create_loss for eval mode and custom loss_fn."""
    loss = np.array([[1.], [2.]], dtype=np.float32)
    logits_input = np.array([[-10.], [10.]], dtype=np.float32)
    labels_input = np.array([[1], [0]], dtype=np.int64)

    def _loss_fn(labels, logits):
      check_labels = tf.debugging.Assert(
          tf.reduce_all(tf.math.equal(labels, labels_input)), data=[labels])
      check_logits = tf.debugging.Assert(
          tf.reduce_all(tf.math.equal(logits, logits_input)), data=[logits])
      with tf.control_dependencies([check_labels, check_logits]):
        return tf.constant(loss)

    head = head_lib.BinaryClassHead(loss_fn=_loss_fn)

    actual_training_loss = head.loss(
        logits=logits_input,
        labels=labels_input,
        features={'x': np.array(((42,),), dtype=np.int32)},
        mode=ModeKeys.EVAL)
    self.assertAllClose(np.sum(loss) / 2., self.evaluate(actual_training_loss))

  def test_eval_create_loss_loss_fn_wrong_shape(self):
    """Tests custom loss_fn that returns Tensor of unexpected shape."""
    loss = np.array([1., 2.], dtype=np.float32)

    def _loss_fn(labels, logits):
      del labels, logits  # Unused
      return tf.constant(loss)

    head = head_lib.BinaryClassHead(loss_fn=_loss_fn)

    logits = np.array([[-10.], [10.]], dtype=np.float32)
    labels = np.array([[1], [0]], dtype=np.int64)
    features = {'x': np.array(((42,),), dtype=np.int32)}

    if tf.executing_eagerly():
      with self.assertRaisesRegexp(ValueError, 'loss_shape'):
        head.loss(
            logits=logits, labels=labels, features=features, mode=ModeKeys.EVAL)
    else:
      actual_training_loss = head.loss(
          logits=logits, labels=labels, features=features, mode=ModeKeys.EVAL)
      with self.assertRaisesRegexp(
          tf.errors.InvalidArgumentError,
          r'\[loss_fn must return Tensor of shape \[D0, D1, ... DN, 1\]\. \] '
          r'\[logits_shape: \] \[2 1\] \[loss_shape: \] \[2\]'):
        with self.cached_session():
          test_lib._initialize_variables(self, tf.compat.v1.train.Scaffold())
          actual_training_loss.eval()

  def test_eval_labels_none(self):
    """Tests that error is raised when labels is None."""
    head = head_lib.BinaryClassHead()

    with self.assertRaisesRegexp(
        ValueError, r'You must provide a labels Tensor\. Given: None\.'):
      head.loss(
          logits=np.array(((45,), (-41,),), dtype=np.float32),
          labels=None,
          features={'x': np.array(((42,),), dtype=np.int32)},
          mode=ModeKeys.EVAL)

  def test_eval(self):
    head = head_lib.BinaryClassHead()
    logits = np.array(((45,), (-41,),), dtype=np.float32)
    labels = np.array(((1,), (1,),), dtype=np.int32)
    features = {'x': np.array(((42,),), dtype=np.int32)}

    # loss = sum(cross_entropy(labels, logits)) / batch_size
    #      = sum(0, 41) / 2 = 41 / 2 = 20.5
    expected_loss = 20.5
    keys = metric_keys.MetricKeys
    expected_metrics = {
        keys.LOSS_MEAN: expected_loss,
        keys.ACCURACY: 1. / 2,
        keys.PRECISION: 1.,
        keys.RECALL: 1. / 2,
        keys.PREDICTION_MEAN: 1. / 2,
        keys.LABEL_MEAN: 2. / 2,
        keys.ACCURACY_BASELINE: 2. / 2,
        keys.AUC: 0.,
        keys.AUC_PR: 1.,
    }
    if tf.executing_eagerly():
      eval_metrics = head.metrics()
      updated_metrics = head.update_metrics(eval_metrics, features, logits,
                                            labels)
      self.assertItemsEqual(expected_metrics.keys(), updated_metrics.keys())
      self.assertAllClose(
          expected_metrics,
          {k: updated_metrics[k].result() for k in updated_metrics})
      loss = head.loss(labels, logits, features=features, mode=ModeKeys.EVAL)
      self.assertIsNotNone(loss)
      self.assertAllClose(expected_loss, loss)
      return

    # Create estimator spec.
    spec = head.create_estimator_spec(
        features=features,
        mode=ModeKeys.EVAL,
        logits=logits,
        labels=labels,
        trainable_variables=[tf.Variable([1.0, 2.0], dtype=tf.dtypes.float32)])
    # Assert spec contains expected tensors.
    self.assertIsNotNone(spec.loss)
    self.assertItemsEqual(expected_metrics.keys(), spec.eval_metric_ops.keys())
    self.assertIsNone(spec.train_op)
    self.assertIsNone(spec.export_outputs)
    test_lib._assert_no_hooks(self, spec)

    # Assert predictions, loss, and metrics.
    with self.cached_session() as sess:
      test_lib._initialize_variables(self, spec.scaffold)
      self.assertIsNone(spec.scaffold.summary_op)
      value_ops = {k: spec.eval_metric_ops[k][0] for k in spec.eval_metric_ops}
      update_ops = {k: spec.eval_metric_ops[k][1] for k in spec.eval_metric_ops}
      loss, _ = sess.run((spec.loss, update_ops))
      self.assertAllClose(expected_loss, loss)
      # Check results of value ops (in `metrics`).
      self.assertAllClose(expected_metrics,
                          {k: value_ops[k].eval() for k in value_ops})

  def test_eval_metric_ops_with_head_name(self):
    head = head_lib.BinaryClassHead(name='some_binary_head')
    logits = np.array(((45,), (-41,),), dtype=np.float32)
    labels = np.array(((1,), (1,),), dtype=np.int32)
    features = {'x': np.array(((42,),), dtype=np.int32)}
    keys = metric_keys.MetricKeys
    expected_metric_keys = [
        '{}/some_binary_head'.format(keys.LOSS_MEAN),
        '{}/some_binary_head'.format(keys.ACCURACY),
        '{}/some_binary_head'.format(keys.PRECISION),
        '{}/some_binary_head'.format(keys.RECALL),
        '{}/some_binary_head'.format(keys.PREDICTION_MEAN),
        '{}/some_binary_head'.format(keys.LABEL_MEAN),
        '{}/some_binary_head'.format(keys.ACCURACY_BASELINE),
        '{}/some_binary_head'.format(keys.AUC),
        '{}/some_binary_head'.format(keys.AUC_PR),
    ]
    eval_metrics = head.metrics()
    updated_metrics = head.update_metrics(eval_metrics, features, logits,
                                          labels)
    self.assertItemsEqual(expected_metric_keys, updated_metrics.keys())

  def test_eval_with_regularization_losses(self):
    head = head_lib.BinaryClassHead()
    logits = np.array(((45,), (-41,),), dtype=np.float32)
    labels = np.array(((1,), (1,),), dtype=np.int32)
    features = {'x': np.array(((42,),), dtype=np.int32)}