# Copyright 2016 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 `models.py` (model cloning, mainly)."""

import tensorflow.compat.v2 as tf

import functools
import os

from absl.testing import parameterized
import numpy as np

import keras
from keras import backend
from keras.testing_infra import test_combinations
from keras import metrics
from keras import models
from keras.optimizers import optimizer_v1
from keras.testing_infra import test_utils


class TestModel(keras.Model):
  """A model subclass."""

  def __init__(self, n_outputs=4, trainable=True):
    """A test class with one dense layer and number of outputs as a variable."""
    super(TestModel, self).__init__()
    self.layer1 = keras.layers.Dense(n_outputs)
    self.n_outputs = tf.Variable(n_outputs, trainable=trainable)

  def call(self, x):
    return self.layer1(x)


def _get_layers(input_shape=(4,), add_input_layer=False):
  if add_input_layer:
    model_layers = [keras.layers.InputLayer(input_shape=input_shape),
                    keras.layers.Dense(4)]
  elif input_shape:
    model_layers = [keras.layers.Dense(4, input_shape=input_shape)]
  else:
    model_layers = [keras.layers.Dense(4)]

  model_layers += [
      keras.layers.BatchNormalization(),
      keras.layers.Dropout(0.5),
      keras.layers.Dense(4)]

  return model_layers


def _get_model(input_shape=(4,)):
  model_layers = _get_layers(input_shape=None, add_input_layer=False)
  return test_utils.get_model_from_layers(
      model_layers, input_shape=input_shape)


class TestModelCloning(test_combinations.TestCase):

  @test_combinations.run_all_keras_modes
  @parameterized.named_parameters([
      {'testcase_name': 'has_input_layer',
       'input_shape': (4,),
       'add_input_layer': True,
       'share_weights': False},
      {'testcase_name': 'no_input_layer',
       'input_shape': None,
       'add_input_layer': False,
       'share_weights': False},
      {'testcase_name': 'has_input_layer_share_weights',
       'input_shape': (4,),
       'add_input_layer': True,
       'share_weights': True},
      {'testcase_name': 'no_input_layer_share_weights',
       'input_shape': None,
       'add_input_layer': False,
       'share_weights': True},
  ])
  def test_clone_sequential_model(
      self, input_shape, add_input_layer, share_weights):

    if share_weights:
      clone_fn = functools.partial(
          keras.models._clone_sequential_model, layer_fn=models.share_weights)
    else:
      clone_fn = keras.models.clone_model

    val_a = np.random.random((10, 4))
    model = models.Sequential(_get_layers(input_shape, add_input_layer))
    # Sanity check
    self.assertEqual(
        isinstance(
            list(model._flatten_layers(include_self=False, recursive=False))[0],
            keras.layers.InputLayer), add_input_layer)
    self.assertEqual(model._is_graph_network, add_input_layer)

    # With placeholder creation -- clone model should have an InputLayer
    # if the original model has one.
    new_model = clone_fn(model)
    self.assertEqual(
        isinstance(
            list(
                new_model._flatten_layers(include_self=False,
                                          recursive=False))[0],
            keras.layers.InputLayer), add_input_layer)
    self.assertEqual(new_model._is_graph_network, model._is_graph_network)
    if input_shape and not tf.compat.v1.executing_eagerly_outside_functions():
      # update ops from batch norm needs to be included
      self.assertGreaterEqual(len(new_model.updates), 2)

    # On top of new tensor  -- clone model should always have an InputLayer.
    input_a = keras.Input(shape=(4,), name='a')
    new_model = clone_fn(model, input_tensors=input_a)
    self.assertIsInstance(
        list(new_model._flatten_layers(include_self=False, recursive=False))[0],
        keras.layers.InputLayer)
    # The new models inputs should have the properties of the new input tensor
    if tf.__internal__.tf2.enabled():
      # In TF1, the new model will be a:0
      self.assertEqual(new_model.input_names[0], input_a.name)
    self.assertEqual(new_model.inputs[0].shape, input_a.shape)
    self.assertTrue(new_model._is_graph_network)

    # On top of new, non-Keras tensor  -- clone model should always have an
    # InputLayer.
    if not tf.executing_eagerly():
      # TODO(b/121277734):Skip Eager contexts, as Input() layers raise an error
      # saying they should not be used with EagerTensors
      input_a = keras.backend.variable(val_a)
      new_model = clone_fn(model, input_tensors=input_a)
      self.assertIsInstance(
          list(new_model._flatten_layers(include_self=False,
                                         recursive=False))[0],
          keras.layers.InputLayer)
      self.assertTrue(new_model._is_graph_network)

  @test_combinations.run_all_keras_modes
  @parameterized.named_parameters([
      {'testcase_name': 'clone_weights', 'share_weights': False},
      {'testcase_name': 'share_weights', 'share_weights': True},
  ])
  def test_clone_functional_model(self, share_weights):