# 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.
# ==============================================================================
"""Test the keras ResNet model with ImageNet data."""

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

from absl.testing import parameterized
import tensorflow as tf

from tensorflow.python.eager import context
from official.utils.misc import keras_utils
from official.utils.testing import integration
from official.vision.image_classification import imagenet_preprocessing
from official.vision.image_classification import resnet_imagenet_main


@parameterized.parameters(
    "resnet",
    "resnet_polynomial_decay",
    "mobilenet",
    "mobilenet_polynomial_decay")
class KerasImagenetTest(tf.test.TestCase):
  """Unit tests for Keras Models with ImageNet."""
  _default_flags_dict = [
      "-batch_size", "4",
      "-train_steps", "1",
      "-use_synthetic_data", "true",
      "-data_format", "channels_last",
  ]
  _extra_flags_dict = {
      "resnet": [
          "-model", "resnet50_v1.5",
          "-optimizer", "resnet50_default",
      ],
      "resnet_polynomial_decay": [
          "-model", "resnet50_v1.5",
          "-optimizer", "resnet50_default",
          "-pruning_method", "polynomial_decay",
          "-use_tf_keras_layers", "true",
      ],
      "mobilenet": [
          "-model", "mobilenet",
          "-optimizer", "mobilenet_default",
      ],
      "mobilenet_polynomial_decay": [
          "-model", "mobilenet",
          "-optimizer", "mobilenet_default",
          "-pruning_method", "polynomial_decay",
      ],
  }
  _tempdir = None

  @classmethod
  def setUpClass(cls):  # pylint: disable=invalid-name
    super(KerasImagenetTest, cls).setUpClass()
    resnet_imagenet_main.define_imagenet_keras_flags()

  def setUp(self):
    super(KerasImagenetTest, self).setUp()
    imagenet_preprocessing.NUM_IMAGES["validation"] = 4
    self.policy = \
        tf.compat.v2.keras.mixed_precision.experimental.global_policy()

  def tearDown(self):
    super(KerasImagenetTest, self).tearDown()
    tf.io.gfile.rmtree(self.get_temp_dir())
    tf.compat.v2.keras.mixed_precision.experimental.set_policy(self.policy)

  def get_extra_flags_dict(self, flags_key):
    return self._extra_flags_dict[flags_key] + self._default_flags_dict

  def test_end_to_end_no_dist_strat(self, flags_key):
    """Test Keras model with 1 GPU, no distribution strategy."""
    config = keras_utils.get_config_proto_v1()
    tf.compat.v1.enable_eager_execution(config=config)

    extra_flags = [
        "-distribution_strategy", "off",
    ]
    extra_flags = extra_flags + self.get_extra_flags_dict(flags_key)

    integration.run_synthetic(
        main=resnet_imagenet_main.run,
        tmp_root=self.get_temp_dir(),
        extra_flags=extra_flags
    )

  def test_end_to_end_graph_no_dist_strat(self, flags_key):
    """Test Keras model in legacy graph mode with 1 GPU, no dist strat."""
    extra_flags = [
        "-enable_eager", "false",
        "-distribution_strategy", "off",
    ]
    extra_flags = extra_flags + self.get_extra_flags_dict(flags_key)

    integration.run_synthetic(
        main=resnet_imagenet_main.run,
        tmp_root=self.get_temp_dir(),
        extra_flags=extra_flags
    )

  def test_end_to_end_1_gpu(self, flags_key):
    """Test Keras model with 1 GPU."""
    config = keras_utils.get_config_proto_v1()
    tf.compat.v1.enable_eager_execution(config=config)