# 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.
# ==============================================================================
"""Correctness tests for tf.keras DNN model using DistributionStrategy."""

import tensorflow.compat.v2 as tf

import numpy as np

import keras
from keras import backend
from keras import testing_utils
from keras.distribute import keras_correctness_test_base
from keras.distribute import strategy_combinations
from keras.optimizer_v2 import gradient_descent as gradient_descent_keras


def all_strategy_combinations_with_eager_and_graph_modes():
  return (tf.__internal__.test.combinations.combine(
      distribution=strategy_combinations.all_strategies,
      mode=['graph', 'eager']) + tf.__internal__.test.combinations.combine(
          distribution=strategy_combinations.multi_worker_mirrored_strategies,
          mode='eager'))


def all_strategy_combinations_with_graph_mode():
  return (tf.__internal__.test.combinations.combine(
      distribution=keras_correctness_test_base.all_strategies,
      mode=['graph']))


def is_default_strategy(strategy):
  with strategy.scope():
    return not tf.distribute.has_strategy()


@testing_utils.run_all_without_tensor_float_32(
    'Uses Dense layers, which call matmul')
class TestDistributionStrategyDnnCorrectness(
    keras_correctness_test_base.TestDistributionStrategyCorrectnessBase):

  def get_model(self,
                initial_weights=None,
                distribution=None,
                input_shapes=None):
    with keras_correctness_test_base.MaybeDistributionScope(distribution):
      # We add few non-linear layers to make it non-trivial.
      model = keras.Sequential()
      model.add(keras.layers.Dense(10, activation='relu', input_shape=(1,)))
      model.add(
          keras.layers.Dense(
              10,
              activation='relu',
              kernel_regularizer=keras.regularizers.l2(1e-4)))
      model.add(keras.layers.Dense(10, activation='relu'))
      model.add(keras.layers.Dense(1))

      if initial_weights:
        model.set_weights(initial_weights)

      model.compile(
          loss=keras.losses.mean_squared_error,
          optimizer=gradient_descent_keras.SGD(0.05),
          metrics=['mse'])
      return model

  def get_data(self):
    x_train = np.random.rand(9984, 1).astype('float32')
    y_train = 3 * x_train
    x_predict = np.array([[1.], [2.], [3.], [4.]], dtype=np.float32)
    return x_train, y_train, x_predict

  def get_data_with_partial_last_batch(self):
    x_train = np.random.rand(10000, 1).astype('float32')
    y_train = 3 * x_train
    x_eval = np.random.rand(10000, 1).astype('float32')
    y_eval = 3 * x_eval
    x_predict = np.array([[1.], [2.], [3.], [4.]], dtype=np.float32)
    return x_train, y_train, x_eval, y_eval, x_predict

  def get_data_with_partial_last_batch_eval(self):
    x_train = np.random.rand(9984, 1).astype('float32')
    y_train = 3 * x_train
    x_eval = np.random.rand(10000, 1).astype('float32')
    y_eval = 3 * x_eval
    x_predict = np.array([[1.], [2.], [3.], [4.]], dtype=np.float32)
    return x_train, y_train, x_eval, y_eval, x_predict

  @tf.__internal__.distribute.combinations.generate(
      keras_correctness_test_base.all_strategy_and_input_config_combinations() +
      keras_correctness_test_base.multi_worker_mirrored_eager())
  def test_dnn_correctness(self, distribution, use_numpy, use_validation_data):
    self.run_correctness_test(distribution, use_numpy, use_validation_data)

  @tf.__internal__.distribute.combinations.generate(
      keras_correctness_test_base
      .test_combinations_with_tpu_strategies_graph() +
      keras_correctness_test_base.multi_worker_mirrored_eager())
  def test_dnn_correctness_with_partial_last_batch_eval(self, distribution,
                                                        use_numpy,
                                                        use_validation_data):
    self.run_correctness_test(
        distribution, use_numpy, use_validation_data, partial_last_batch='eval')

  @tf.__internal__.distribute.combinations.generate(
      keras_correctness_test_base
      .strategy_minus_tpu_and_input_config_combinations_eager() +
      keras_correctness_test_base.multi_worker_mirrored_eager())
  def test_dnn_correctness_with_partial_last_batch(self, distribution,
                                                   use_numpy,
                                                   use_validation_data):
    distribution.extended.experimental_enable_get_next_as_optional = True
    self.run_correctness_test(
        distribution,
        use_numpy,
        use_validation_data,
        partial_last_batch='train_and_eval',
        training_epochs=1)

  @tf.__internal__.distribute.combinations.generate(all_strategy_combinations_with_graph_mode())
  def test_dnn_with_dynamic_learning_rate(self, distribution):
    self.run_dynamic_lr_test(distribution)


class TestDistributionStrategyDnnMetricCorrectness(
    keras_correctness_test_base.TestDistributionStrategyCorrectnessBase):

  def get_model(self,
                distribution=None,
                input_shapes=None):
    with distribution.scope():
      model = keras.Sequential()
      model.add(
          keras.layers.Dense(1, input_shape=(1,), kernel_initializer='ones'))
      model.compile(
          loss=keras.losses.mean_squared_error,
          optimizer=gradient_descent_keras.SGD(0.05),
          metrics=[keras.metrics.BinaryAccuracy()])
    return model

  def run_metric_correctness_test(self, distribution):
    with self.cached_session():
      self.set_up_test_config()

      x_train, y_train, _ = self.get_data()
      model = self.get_model(
          distribution=distribution)

      batch_size = 64
      batch_size = (
          keras_correctness_test_base.get_batch_size(batch_size, distribution))
      train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train))
      train_dataset = (
          keras_correctness_test_base.batch_wrapper(train_dataset, batch_size))

      history = model.fit(x=train_dataset, epochs=2, steps_per_epoch=10)
      self.assertEqual(history.history['binary_accuracy'], [1.0, 1.0])

  @tf.__internal__.distribute.combinations.generate(
      all_strategy_combinations_with_eager_and_graph_modes())
  def test_simple_dnn_metric_correctness(self, distribution):
    self.run_metric_correctness_test(distribution)


class TestDistributionStrategyDnnMetricEvalCorrectness(
    keras_correctness_test_base.TestDistributionStrategyCorrectnessBase):

  def get_model(self,
                distribution=None,
                input_shapes=None):
    with distribution.scope():
      model = keras.Sequential()
      model.add(
          keras.layers.Dense(
              3, activation='relu', input_dim=4, kernel_initializer='ones'))
      model.add(
          keras.layers.Dense(
              1, activation='sigmoid', kernel_initializer='ones'))
      model.compile(
          loss='mae',
          metrics=['accuracy', keras.metrics.BinaryAccuracy()],
          optimizer=tf.compat.v1.train.GradientDescentOptimizer(0.001))
    return model

  def run_eval_metrics_correctness_test(self, distribution):
    with self.cached_session():
      self.set_up_test_config()

      model = self.get_model(
          distribution=distribution)

      # verify correctness of stateful and stateless metrics.
      x = np.ones((100, 4)).astype('float32')
      y = np.ones((100, 1)).astype('float32')
      dataset = tf.data.Dataset.from_tensor_slices((x, y)).repeat()
      dataset = keras_correctness_test_base.batch_wrapper(dataset, 4)
      outs = model.evaluate(dataset, steps=10)
      self.assertEqual(outs[1], 1.)
      self.assertEqual(outs[2], 1.)

      y = np.zeros((100, 1)).astype('float32')
      dataset = tf.data.Dataset.from_tensor_slices((x, y)).repeat()
      dataset = keras_correctness_test_base.batch_wrapper(dataset, 4)
      outs = model.evaluate(dataset, steps=10)
      self.assertEqual(outs[1], 0.)
      self.assertEqual(outs[2], 0.)

  @tf.__internal__.distribute.combinations.generate(
      all_strategy_combinations_with_eager_and_graph_modes())
  def test_identity_model_metric_eval_correctness(self, distribution):
    self.run_eval_metrics_correctness_test(distribution)


class SubclassedModel(keras.Model):

  def __init__(self, initial_weights, input_shapes):
    super(SubclassedModel, self).__init__()
    self.dense1 = keras.layers.Dense(10, activation='relu', input_shape=(1,))
    self.dense2 = keras.layers.Dense(
        10, activation='relu', kernel_regularizer=keras.regularizers.l2(1e-4))
    self.dense3 = keras.layers.Dense(10, activation='relu')
    self.dense4 = keras.layers.Dense(1)
    if input_shapes:
      self.build(input_shapes)
    else:
      # This covers cases when the input is DatasetV1Adapter.
      self.build((None, 1))
    if initial_weights:
      self.set_weights(initial_weights)

  def call(self, inputs):
    x = self.dense1(inputs)
    x = self.dense2(x)
    x = self.dense3(x)
    return self.dense4(x)


@testing_utils.run_all_without_tensor_float_32(
    'Uses Dense layers, which call matmul')
class TestDistributionStrategyDnnCorrectnessWithSubclassedModel(
    TestDistributionStrategyDnnCorrectness):

  def get_model(self,
                initial_weights=None,
                distribution=None,
                input_shapes=None):
    with keras_correctness_test_base.MaybeDistributionScope(distribution):
      model = SubclassedModel(initial_weights, input_shapes)

      model.compile(
          loss=keras.losses.mean_squared_error,
          optimizer=gradient_descent_keras.SGD(0.05),
          metrics=['mse'])
      return model

  @tf.__internal__.distribute.combinations.generate(
      keras_correctness_test_base.all_strategy_and_input_config_combinations() +
      keras_correctness_test_base.multi_worker_mirrored_eager())
  def test_dnn_correctness(self, distribution, use_numpy, use_validation_data):
    if (tf.executing_eagerly()) or is_default_strategy(distribution):
      self.run_correctness_test(distribution, use_numpy, use_validation_data)
    elif (backend.is_tpu_strategy(distribution)
          and not tf.executing_eagerly()):
      with self.assertRaisesRegex(
          ValueError,
          'Expected `model` argument to be a functional `Model` instance, '
          'but got a subclass model instead.'):
        self.run_correctness_test(distribution, use_numpy, use_validation_data)
    else:
      with self.assertRaisesRegex(
          ValueError,
          'We currently do not support distribution strategy with a '
          '`Sequential` model that is created without `input_shape`/'
          '`input_dim` set in its first layer or a subclassed model.'):
        self.run_correctness_test(distribution, use_numpy, use_validation_data)

  @tf.__internal__.distribute.combinations.generate(all_strategy_combinations_with_graph_mode())
  def test_dnn_with_dynamic_learning_rate(self, distribution):
    if ((tf.executing_eagerly()
         and not backend.is_tpu_strategy(distribution))
        or is_default_strategy(distribution)):
      self.run_dynamic_lr_test(distribution)
    elif backend.is_tpu_strategy(distribution):
      with self.assertRaisesRegex(
          ValueError,
          'Expected `model` argument to be a functional `Model` instance, '
          'but got a subclass model instead.'):
        self.run_dynamic_lr_test(distribution)
    else:
      with self.assertRaisesRegex(
          ValueError,
          'We currently do not support distribution strategy with a '
          '`Sequential` model that is created without `input_shape`/'
          '`input_dim` set in its first layer or a subclassed model.'):
        self.run_dynamic_lr_test(distribution)

  @tf.__internal__.distribute.combinations.generate(
      keras_correctness_test_base.test_combinations_with_tpu_strategies_graph())
  def test_dnn_correctness_with_partial_last_batch_eval(self, distribution,
                                                        use_numpy,
                                                        use_validation_data):
    with self.assertRaisesRegex(
        ValueError,
        'Expected `model` argument to be a functional `Model` instance, '
        'but got a subclass model instead.'):
      self.run_correctness_test(
          distribution,
          use_numpy,
          use_validation_data,
          partial_last_batch='eval')


if __name__ == '__main__':
  tf.__internal__.distribute.multi_process_runner.test_main()
