# Copyright 2019 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 custom training loops."""

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

from absl.testing import parameterized
import numpy as np

from tensorflow.python import keras
from tensorflow.python.eager import backprop
from tensorflow.python.eager import def_function
from tensorflow.python.framework import ops
from tensorflow.python.keras import keras_parameterized
from tensorflow.python.keras import testing_utils
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.platform import test


class LayerWithLosses(keras.layers.Layer):

  def build(self, input_shape):
    self.v = self.add_weight(
        name='hey',
        shape=(),
        initializer='ones',
        regularizer=keras.regularizers.l1(100))

  def call(self, inputs):
    self.add_loss(math_ops.reduce_sum(inputs))
    return self.v * inputs


class LayerWithMetrics(keras.layers.Layer):

  def build(self, input_shape):
    self.mean = keras.metrics.Mean(name='mean_object')

  def call(self, inputs):
    self.add_metric(
        math_ops.reduce_mean(inputs), name='mean_tensor', aggregation='mean')
    self.add_metric(self.mean(inputs))
    return inputs


def add_loss_step(defun):
  optimizer = keras.optimizer_v2.adam.Adam()
  model = testing_utils.get_model_from_layers([LayerWithLosses()],
                                              input_shape=(10,))

  def train_step(x):
    with backprop.GradientTape() as tape:
      model(x)
      assert len(model.losses) == 2
      loss = math_ops.reduce_sum(model.losses)
    gradients = tape.gradient(loss, model.trainable_weights)
    optimizer.apply_gradients(zip(gradients, model.trainable_weights))
    return loss

  if defun:
    train_step = def_function.function(train_step)

  x = array_ops.ones((10, 10))
  return train_step(x)


def batch_norm_step(defun):
  optimizer = keras.optimizer_v2.adadelta.Adadelta()
  model = testing_utils.get_model_from_layers([
      keras.layers.BatchNormalization(momentum=0.9),
      keras.layers.Dense(1, kernel_initializer='zeros', activation='softmax')
  ],
                                              input_shape=(10,))

  def train_step(x, y):
    with backprop.GradientTape() as tape:
      y_pred = model(x, training=True)
      loss = keras.losses.binary_crossentropy(y, y_pred)
    gradients = tape.gradient(loss, model.trainable_weights)
    optimizer.apply_gradients(zip(gradients, model.trainable_weights))
    return loss, model(x, training=False)

  if defun:
    train_step = def_function.function(train_step)

  x, y = array_ops.ones((10, 10)), array_ops.ones((10, 1))
  return train_step(x, y)


def add_metric_step(defun):
  optimizer = keras.optimizer_v2.rmsprop.RMSprop()
  model = testing_utils.get_model_from_layers([
      LayerWithMetrics(),
      keras.layers.Dense(1, kernel_initializer='zeros', activation='softmax')
  ],
                                              input_shape=(10,))

  def train_step(x, y):
    with backprop.GradientTape() as tape:
      y_pred_1 = model(x)
      y_pred_2 = model(2 * x)
      y_pred = y_pred_1 + y_pred_2
      loss = keras.losses.mean_squared_error(y, y_pred)
    gradients = tape.gradient(loss, model.trainable_weights)
    optimizer.apply_gradients(zip(gradients, model.trainable_weights))
    assert len(model.metrics) == 2
    return [m.result() for m in model.metrics]

  if defun:
    train_step = def_function.function(train_step)

  x, y = array_ops.ones((10, 10)), array_ops.zeros((10, 1))
  metrics = train_step(x, y)
  assert np.allclose(metrics[0], 1.5)
  assert np.allclose(metrics[1], 1.5)
  return metrics


@keras_parameterized.run_with_all_model_types
class CustomTrainingLoopTest(keras_parameterized.TestCase):

  @parameterized.named_parameters(('add_loss_step', add_loss_step),
                                  ('add_metric_step', add_metric_step),
                                  ('batch_norm_step', batch_norm_step))
  def test_eager_and_tf_function(self, train_step):
    eager_result = train_step(defun=False)
    fn_result = train_step(defun=True)
    self.assertAllClose(eager_result, fn_result)


if __name__ == '__main__':
  ops.enable_eager_execution()
  test.main()
