# Copyright 2017 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 BaselineEstimator."""

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

import os
import shutil
import tempfile

import numpy as np
import six
import tensorflow as tf
from tensorflow.python.framework import ops
from tensorflow_estimator.python.estimator.canned import baseline
from tensorflow_estimator.python.estimator.canned import metric_keys
from tensorflow_estimator.python.estimator.export import export
from tensorflow_estimator.python.estimator.head import regression_head
from tensorflow_estimator.python.estimator.inputs import numpy_io

# Names of variables created by model.
BIAS_NAME = 'baseline/bias'


def assert_close(expected, actual, rtol=1e-04, name='assert_close'):
  with ops.name_scope(name, 'assert_close', (expected, actual, rtol)) as scope:
    expected = ops.convert_to_tensor(expected, name='expected')
    actual = ops.convert_to_tensor(actual, name='actual')
    rdiff = tf.math.abs(expected - actual, 'diff') / tf.math.abs(expected)
    rtol = ops.convert_to_tensor(rtol, name='rtol')
    return tf.compat.v1.debugging.assert_less(
        rdiff,
        rtol,
        data=('Condition expected =~ actual did not hold element-wise:'
              'expected = ', expected, 'actual = ', actual, 'rdiff = ', rdiff,
              'rtol = ', rtol,),
        name=scope)


def save_variables_to_ckpt(model_dir):
  init_all_op = [tf.compat.v1.initializers.global_variables()]
  with tf.compat.v1.Session() as sess:
    sess.run(init_all_op)
    tf.compat.v1.train.Saver().save(sess, os.path.join(model_dir, 'model.ckpt'))


def _baseline_estimator_fn(weight_column=None, label_dimension=1, **kwargs):
  return baseline.BaselineEstimatorV2(
      head=regression_head.RegressionHead(
          weight_column=weight_column,
          label_dimension=label_dimension,
          loss_reduction=tf.losses.Reduction.SUM_OVER_BATCH_SIZE),
      **kwargs)


def mock_optimizer_v2(testcase, expected_loss=None):
  """Creates a mock optimizer to test the train method.

  Args:
    testcase: A TestCase instance.
    expected_loss: If given, will assert the loss value.

  Returns:
    A mock Optimizer.
  """
  expected_var_names = ['%s:0' % BIAS_NAME]

  class _Optimizer(tf.keras.optimizers.Optimizer):

    def get_updates(self, loss, params):
      trainable_vars = params
      testcase.assertItemsEqual(expected_var_names,
                                [var.name for var in trainable_vars])

      # Verify loss. We can't check the value directly, so we add an assert op.
      testcase.assertEquals(0, loss.shape.ndims)
      if expected_loss is None:
        if self.iterations is not None:
          return [self.iterations.assign_add(1).op]
        return [tf.no_op()]
      assert_loss = assert_close(
          tf.cast(expected_loss, name='expected', dtype=tf.dtypes.float32),
          loss,
          name='assert_loss')
      with tf.control_dependencies((assert_loss,)):
        if self.iterations is not None:
          return [self.iterations.assign_add(1).op]
        return [tf.no_op()]

    def get_config(self):
      config = super(_Optimizer, self).get_config()
      return config

  optimizer = _Optimizer(name='my_optimizer')

  return optimizer


class BaselineEstimatorEvaluationTest(tf.test.TestCase):

  def setUp(self):
    self._model_dir = tempfile.mkdtemp()

  def tearDown(self):
    if self._model_dir:
      tf.compat.v1.summary.FileWriterCache.clear()
      shutil.rmtree(self._model_dir)

  def test_evaluation_batch(self):
    """Tests evaluation for batch_size==2."""
    with tf.Graph().as_default():
      tf.Variable([13.0], name=BIAS_NAME)
      tf.Variable(
          100, name=tf.compat.v1.GraphKeys.GLOBAL_STEP, dtype=tf.dtypes.int64)
      save_variables_to_ckpt(self._model_dir)

    baseline_estimator = _baseline_estimator_fn(model_dir=self._model_dir)
    eval_metrics = baseline_estimator.evaluate(
        input_fn=lambda: ({'age': ((1,), (1,))}, ((10.,), (10.,))), steps=1)

    # Logit is bias = 13, while label is 10.
    # Loss per example is 3**2 = 9.
    # Training loss is the sum over batch size = (9 + 9) / 2 = 9
    # Average loss is the average over batch = 9
    self.assertDictEqual(
        {
            metric_keys.MetricKeys.LOSS: 9.,
            metric_keys.MetricKeys.LOSS_MEAN: 9.,
            metric_keys.MetricKeys.PREDICTION_MEAN: 13.,
            metric_keys.MetricKeys.LABEL_MEAN: 10.,
            tf.compat.v1.GraphKeys.GLOBAL_STEP: 100
        }, eval_metrics)

  def test_evaluation_weights(self):
    """Tests evaluation with weights."""
    with tf.Graph().as_default():
      tf.Variable([13.0], name=BIAS_NAME)
      tf.Variable(
          100, name=tf.compat.v1.GraphKeys.GLOBAL_STEP, dtype=tf.dtypes.int64)
      save_variables_to_ckpt(self._model_dir)

    def _input_fn():
      features = {'age': ((1,), (1,)), 'weights': ((1.,), (2.,))}
      labels = ((10.,), (10.,))
      return features, labels

    baseline_estimator = _baseline_estimator_fn(
        weight_column='weights', model_dir=self._model_dir)
    eval_metrics = baseline_estimator.evaluate(input_fn=_input_fn, steps=1)

    # Logit is bias = 13, while label is 10.
    # Loss per example is 3**2 = 9.
    # Training loss is the weighted sum over batch size= (9 + 2*9) / 2 = 13.5
    # average loss is the weighted average = 9 + 2*9 / (1 + 2) = 9
    self.assertDictEqual(
        {
            metric_keys.MetricKeys.LOSS: 13.5,
            metric_keys.MetricKeys.LOSS_MEAN: 9.,
            metric_keys.MetricKeys.PREDICTION_MEAN: 13.,
            metric_keys.MetricKeys.LABEL_MEAN: 10.,
            tf.compat.v1.GraphKeys.GLOBAL_STEP: 100
        }, eval_metrics)

  def test_evaluation_for_multi_dimensions(self):
    label_dim = 2
    with tf.Graph().as_default():
      tf.Variable([46.0, 58.0], name=BIAS_NAME)
      tf.Variable(100, name='global_step', dtype=tf.dtypes.int64)
      save_variables_to_ckpt(self._model_dir)

    baseline_estimator = _baseline_estimator_fn(
        label_dimension=label_dim, model_dir=self._model_dir)
    input_fn = numpy_io.numpy_input_fn(
        x={
            'age': np.array([[2., 4., 5.]]),
        },
        y=np.array([[46., 58.]]),
        batch_size=1,
        num_epochs=None,
        shuffle=False)
    eval_metrics = baseline_estimator.evaluate(input_fn=input_fn, steps=1)

    self.assertItemsEqual(
        (metric_keys.MetricKeys.LOSS, metric_keys.MetricKeys.LOSS_MEAN,
         metric_keys.MetricKeys.PREDICTION_MEAN,
         metric_keys.MetricKeys.LABEL_MEAN, tf.compat.v1.GraphKeys.GLOBAL_STEP),
        eval_metrics.keys())

    # Logit is bias which is [46, 58]
    self.assertAlmostEqual(0, eval_metrics[metric_keys.MetricKeys.LOSS])


class BaselineEstimatorPredictTest(tf.test.TestCase):

  def setUp(self):
    self._model_dir = tempfile.mkdtemp()

  def tearDown(self):
    if self._model_dir:
      tf.compat.v1.summary.FileWriterCache.clear()
      shutil.rmtree(self._model_dir)

  def test_1d(self):
    """Tests predict when all variables are one-dimensional."""
    with tf.Graph().as_default():
      tf.Variable([.2], name=BIAS_NAME)
      tf.Variable(100, name='global_step', dtype=tf.dtypes.int64)
      save_variables_to_ckpt(self._model_dir)

    baseline_estimator = _baseline_estimator_fn(model_dir=self._model_dir)

    predict_input_fn = numpy_io.numpy_input_fn(
        x={'x': np.array([[2.]])},
        y=None,
        batch_size=1,
        num_epochs=1,
        shuffle=False)
    predictions = baseline_estimator.predict(input_fn=predict_input_fn)
    predicted_scores = list([x['predictions'] for x in predictions])
    # x * weight + bias = 2. * 10. + .2 = 20.2
    self.assertAllClose([[.2]], predicted_scores)

  def testMultiDim(self):
    """Tests predict when all variables are multi-dimenstional."""
    batch_size = 2
    label_dimension = 3
    with tf.Graph().as_default():
      tf.Variable(  # shape=[label_dimension]
          [.2, .4, .6], name=BIAS_NAME)
      tf.Variable(100, name='global_step', dtype=tf.dtypes.int64)
      save_variables_to_ckpt(self._model_dir)

    baseline_estimator = _baseline_estimator_fn(
        label_dimension=label_dimension, model_dir=self._model_dir)

    predict_input_fn = numpy_io.numpy_input_fn(
        # x shape=[batch_size, x_dim]
        x={'x': np.array([[1., 2., 3., 4.], [5., 6., 7., 8.]])},
        y=None,
        batch_size=batch_size,
        num_epochs=1,
        shuffle=False)
    predictions = baseline_estimator.predict(input_fn=predict_input_fn)
    predicted_scores = list([x['predictions'] for x in predictions])
    # score = bias, shape=[batch_size, label_dimension]
    self.assertAllClose([[0.2, 0.4, 0.6], [0.2, 0.4, 0.6]], predicted_scores)


class BaselineEstimatorIntegrationTest(tf.test.TestCase):

  def setUp(self):
    self._model_dir = tempfile.mkdtemp()

  def tearDown(self):
    if self._model_dir:
      tf.compat.v1.summary.FileWriterCache.clear()
      shutil.rmtree(self._model_dir)

  def _test_complete_flow(self, train_input_fn, eval_input_fn, predict_input_fn,
                          input_dimension, label_dimension, prediction_length):
    feature_columns = [
        tf.feature_column.numeric_column('x', shape=(input_dimension,))
    ]
    est = _baseline_estimator_fn(
        label_dimension=label_dimension, model_dir=self._model_dir)

    # TRAIN
    # learn y = x
    est.train(train_input_fn, steps=200)

    # EVALUTE
    scores = est.evaluate(eval_input_fn)
    self.assertEqual(200, scores[tf.compat.v1.GraphKeys.GLOBAL_STEP])
    self.assertIn(metric_keys.MetricKeys.LOSS, six.iterkeys(scores))

    # PREDICT
    predictions = np.array(
        [x['predictions'] for x in est.predict(predict_input_fn)])
    self.assertAllEqual((prediction_length, label_dimension), predictions.shape)

    # EXPORT
    feature_spec = tf.feature_column.make_parse_example_spec(feature_columns)
    serving_input_receiver_fn = export.build_parsing_serving_input_receiver_fn(
        feature_spec)
    export_dir = est.export_saved_model(tempfile.mkdtemp(),
                                        serving_input_receiver_fn)
    self.assertTrue(tf.compat.v1.gfile.Exists(export_dir))

  def test_numpy_input_fn(self):
    """Tests complete flow with numpy_input_fn."""
    label_dimension = 2
    input_dimension = label_dimension
    batch_size = 10
    prediction_length = batch_size
    data = np.linspace(0., 2., batch_size * label_dimension, dtype=np.float32)
    data = data.reshape(batch_size, label_dimension)

    train_input_fn = numpy_io.numpy_input_fn(
        x={'x': data},
        y=data,
        batch_size=batch_size,
        num_epochs=None,
        shuffle=True)
    eval_input_fn = numpy_io.numpy_input_fn(
        x={'x': data},
        y=data,
        batch_size=batch_size,
        num_epochs=1,
        shuffle=False)
    predict_input_fn = numpy_io.numpy_input_fn(
        x={'x': data},
        y=None,
        batch_size=batch_size,
        num_epochs=1,
        shuffle=False)

    self._test_complete_flow(
        train_input_fn=train_input_fn,
        eval_input_fn=eval_input_fn,
        predict_input_fn=predict_input_fn,
        input_dimension=input_dimension,
        label_dimension=label_dimension,
        prediction_length=prediction_length)


class BaselineEstimatorTrainingTest(tf.test.TestCase):

  def setUp(self):
    self._model_dir = tempfile.mkdtemp()

  def tearDown(self):
    if self._model_dir:
      tf.compat.v1.summary.FileWriterCache.clear()
      shutil.rmtree(self._model_dir)

  def _assert_checkpoint(self,
                         label_dimension,
                         expected_global_step,
                         expected_bias=None):
    shapes = {
        name: shape
        for (name, shape) in tf.train.list_variables(self._model_dir)
    }

    self.assertEqual([], shapes[tf.compat.v1.GraphKeys.GLOBAL_STEP])
    self.assertEqual(
        expected_global_step,
        tf.train.load_variable(self._model_dir,
                               tf.compat.v1.GraphKeys.GLOBAL_STEP))

    self.assertEqual([label_dimension], shapes[BIAS_NAME])
    if expected_bias is not None:
      self.assertEqual(expected_bias,
                       tf.train.load_variable(self._model_dir, BIAS_NAME))

  def testFromScratch(self):
    # Create BaselineRegressor.
    label = 5.
    age = 17
    # loss = (logits - label)^2 = (0 - 5.)^2 = 25.
    mock_optimizer = mock_optimizer_v2(self, expected_loss=25.)
    baseline_estimator = _baseline_estimator_fn(
        model_dir=self._model_dir, optimizer=mock_optimizer)

    # Train for a few steps, and validate optimizer and final checkpoint.
    num_steps = 10
    baseline_estimator.train(
        input_fn=lambda: ({'age': ((age,),)}, ((label,),)), steps=num_steps)
    self.assertEqual(
        num_steps,
        baseline_estimator.get_variable_value(mock_optimizer.iterations.name))
    self._assert_checkpoint(
        label_dimension=1, expected_global_step=num_steps, expected_bias=[0.])

  def testFromCheckpoint(self):
    # Create initial checkpoint.
    bias = 7.0
    initial_global_step = 100
    with tf.Graph().as_default():
      tf.Variable([bias], name=BIAS_NAME)
      tf.Variable(
          initial_global_step,
          name=tf.compat.v1.GraphKeys.GLOBAL_STEP,
          dtype=tf.dtypes.int64)
      save_variables_to_ckpt(self._model_dir)

    # logits = bias = 6.
    # loss = (logits - label)^2 = (7 - 5)^2 = 4
    mock_optimizer = mock_optimizer_v2(self, expected_loss=4.)
    baseline_estimator = _baseline_estimator_fn(
        model_dir=self._model_dir, optimizer=mock_optimizer)

    # Train for a few steps, and validate optimizer and final checkpoint.
    num_steps = 10
    baseline_estimator.train(
        input_fn=lambda: ({'age': ((17,),)}, ((5.,),)), steps=num_steps)
    self.assertEqual(
        initial_global_step + num_steps,
        baseline_estimator.get_variable_value(mock_optimizer.iterations.name))
    self._assert_checkpoint(
        label_dimension=1,
        expected_global_step=initial_global_step + num_steps,
        expected_bias=[bias])


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