# 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 gan.python.train."""

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.contrib import layers
from tensorflow.contrib.framework.python.ops import variables as variables_lib
from tensorflow.contrib.gan.python import namedtuples
from tensorflow.contrib.gan.python import train
from tensorflow.contrib.gan.python.features.python import random_tensor_pool
from tensorflow.contrib.slim.python.slim import learning as slim_learning
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.framework import random_seed
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import random_ops
from tensorflow.python.ops import variable_scope
from tensorflow.python.ops import variables
from tensorflow.python.ops.distributions import categorical
from tensorflow.python.platform import test
from tensorflow.python.training import basic_session_run_hooks
from tensorflow.python.training import coordinator
from tensorflow.python.training import gradient_descent
from tensorflow.python.training import sync_replicas_optimizer
from tensorflow.python.training import training_util


def generator_model(inputs):
  return variable_scope.get_variable('dummy_g', initializer=2.0) * inputs


class Generator(object):

  def __call__(self, inputs):
    return generator_model(inputs)


def infogan_generator_model(inputs):
  return variable_scope.get_variable('dummy_g', initializer=2.0) * inputs[0]


class InfoGANGenerator(object):

  def __call__(self, inputs):
    return infogan_generator_model(inputs)


def discriminator_model(inputs, _):
  return variable_scope.get_variable('dummy_d', initializer=2.0) * inputs


class Discriminator(object):

  def __call__(self, inputs, _):
    return discriminator_model(inputs, _)


def infogan_discriminator_model(inputs, _):
  return (variable_scope.get_variable('dummy_d', initializer=2.0) * inputs,
          [categorical.Categorical([1.0])])


class InfoGANDiscriminator(object):

  def __call__(self, inputs, _):
    return infogan_discriminator_model(inputs, _)


def acgan_discriminator_model(inputs, _, num_classes=10):
  return (
      discriminator_model(inputs, _),
      array_ops.one_hot(
          # TODO(haeusser): infer batch size from input
          random_ops.random_uniform(
              [3], maxval=num_classes, dtype=dtypes.int32),
          num_classes))


class ACGANDiscriminator(object):

  def __call__(self, inputs, _, num_classes=10):
    return (
        discriminator_model(inputs, _),
        array_ops.one_hot(
            # TODO(haeusser): infer batch size from input
            random_ops.random_uniform(
                [3], maxval=num_classes, dtype=dtypes.int32),
            num_classes))


def stargan_generator_model(inputs, _):
  """Dummy generator for StarGAN."""

  return variable_scope.get_variable('dummy_g', initializer=0.5) * inputs


class StarGANGenerator(object):

  def __call__(self, inputs, _):
    return stargan_generator_model(inputs, _)


def stargan_discriminator_model(inputs, num_domains):
  """Differentiable dummy discriminator for StarGAN."""

  hidden = layers.flatten(inputs)

  output_src = math_ops.reduce_mean(hidden, axis=1)

  output_cls = layers.fully_connected(
      inputs=hidden,
      num_outputs=num_domains,
      activation_fn=None,
      normalizer_fn=None,
      biases_initializer=None)
  return output_src, output_cls


class StarGANDiscriminator(object):

  def __call__(self, inputs, num_domains):
    return stargan_discriminator_model(inputs, num_domains)


def get_gan_model():
  # TODO(joelshor): Find a better way of creating a variable scope.
  with variable_scope.variable_scope('generator') as gen_scope:
    pass
  with variable_scope.variable_scope('discriminator') as dis_scope:
    pass
  return namedtuples.GANModel(
      generator_inputs=None,
      generated_data=None,
      generator_variables=None,
      generator_scope=gen_scope,
      generator_fn=generator_model,
      real_data=array_ops.ones([1, 2, 3]),
      discriminator_real_outputs=array_ops.ones([1, 2, 3]),
      discriminator_gen_outputs=array_ops.ones([1, 2, 3]),
      discriminator_variables=None,
      discriminator_scope=dis_scope,
      discriminator_fn=discriminator_model)


def get_callable_gan_model():
  ganmodel = get_gan_model()
  return ganmodel._replace(
      generator_fn=Generator(), discriminator_fn=Discriminator())


def create_gan_model():
  return train.gan_model(
      generator_model,
      discriminator_model,
      real_data=array_ops.zeros([1, 2]),
      generator_inputs=random_ops.random_normal([1, 2]))


def create_callable_gan_model():
  return train.gan_model(
      Generator(),
      Discriminator(),
      real_data=array_ops.zeros([1, 2]),
      generator_inputs=random_ops.random_normal([1, 2]))


def get_infogan_model():
  return namedtuples.InfoGANModel(
      *get_gan_model(),
      structured_generator_inputs=[constant_op.constant(0)],
      predicted_distributions=[categorical.Categorical([1.0])],
      discriminator_and_aux_fn=infogan_discriminator_model)


def get_callable_infogan_model():
  return namedtuples.InfoGANModel(
      *get_callable_gan_model(),
      structured_generator_inputs=[constant_op.constant(0)],
      predicted_distributions=[categorical.Categorical([1.0])],
      discriminator_and_aux_fn=infogan_discriminator_model)


def create_infogan_model():
  return train.infogan_model(
      infogan_generator_model,
      infogan_discriminator_model,
      real_data=array_ops.zeros([1, 2]),
      unstructured_generator_inputs=[],
      structured_generator_inputs=[random_ops.random_normal([1, 2])])


def create_callable_infogan_model():
  return train.infogan_model(
      InfoGANGenerator(),
      InfoGANDiscriminator(),
      real_data=array_ops.zeros([1, 2]),
      unstructured_generator_inputs=[],
      structured_generator_inputs=[random_ops.random_normal([1, 2])])


def get_acgan_model():
  return namedtuples.ACGANModel(
      *get_gan_model(),
      one_hot_labels=array_ops.one_hot([0, 1, 2], 10),
      discriminator_real_classification_logits=array_ops.one_hot([0, 1, 3], 10),
      discriminator_gen_classification_logits=array_ops.one_hot([0, 1, 4], 10))


def get_callable_acgan_model():
  return namedtuples.ACGANModel(
      *get_callable_gan_model(),
      one_hot_labels=array_ops.one_hot([0, 1, 2], 10),
      discriminator_real_classification_logits=array_ops.one_hot([0, 1, 3], 10),
      discriminator_gen_classification_logits=array_ops.one_hot([0, 1, 4], 10))


def create_acgan_model():
  return train.acgan_model(
      generator_model,
      acgan_discriminator_model,
      real_data=array_ops.zeros([1, 2]),
      generator_inputs=random_ops.random_normal([1, 2]),
      one_hot_labels=array_ops.one_hot([0, 1, 2], 10))


def create_callable_acgan_model():
  return train.acgan_model(
      Generator(),
      ACGANDiscriminator(),
      real_data=array_ops.zeros([1, 2]),
      generator_inputs=random_ops.random_normal([1, 2]),
      one_hot_labels=array_ops.one_hot([0, 1, 2], 10))


def get_cyclegan_model():
  return namedtuples.CycleGANModel(
      model_x2y=get_gan_model(),
      model_y2x=get_gan_model(),
      reconstructed_x=array_ops.ones([1, 2, 3]),
      reconstructed_y=array_ops.zeros([1, 2, 3]))


def get_callable_cyclegan_model():
  return namedtuples.CycleGANModel(
      model_x2y=get_callable_gan_model(),
      model_y2x=get_callable_gan_model(),
      reconstructed_x=array_ops.ones([1, 2, 3]),
      reconstructed_y=array_ops.zeros([1, 2, 3]))


def create_cyclegan_model():
  return train.cyclegan_model(
      generator_model,
      discriminator_model,
      data_x=array_ops.zeros([1, 2]),
      data_y=array_ops.ones([1, 2]))


def create_callable_cyclegan_model():
  return train.cyclegan_model(
      Generator(),
      Discriminator(),
      data_x=array_ops.zeros([1, 2]),
      data_y=array_ops.ones([1, 2]))


def get_stargan_model():
  """Similar to get_gan_model()."""
  # TODO(joelshor): Find a better way of creating a variable scope.
  with variable_scope.variable_scope('generator') as gen_scope:
    pass
  with variable_scope.variable_scope('discriminator') as dis_scope:
    pass
  return namedtuples.StarGANModel(
      input_data=array_ops.ones([1, 2, 2, 3]),
      input_data_domain_label=array_ops.ones([1, 2]),
      generated_data=array_ops.ones([1, 2, 2, 3]),
      generated_data_domain_target=array_ops.ones([1, 2]),
      reconstructed_data=array_ops.ones([1, 2, 2, 3]),
      discriminator_input_data_source_predication=array_ops.ones([1]),
      discriminator_generated_data_source_predication=array_ops.ones([1]),
      discriminator_input_data_domain_predication=array_ops.ones([1, 2]),
      discriminator_generated_data_domain_predication=array_ops.ones([1, 2]),
      generator_variables=None,
      generator_scope=gen_scope,
      generator_fn=stargan_generator_model,
      discriminator_variables=None,
      discriminator_scope=dis_scope,
      discriminator_fn=stargan_discriminator_model)


def get_callable_stargan_model():
  model = get_stargan_model()
  return model._replace(
      generator_fn=StarGANGenerator(), discriminator_fn=StarGANDiscriminator())


def create_stargan_model():
  return train.stargan_model(
      stargan_generator_model, stargan_discriminator_model,
      array_ops.ones([1, 2, 2, 3]), array_ops.ones([1, 2]))


def create_callable_stargan_model():
  return train.stargan_model(StarGANGenerator(), StarGANDiscriminator(),
                             array_ops.ones([1, 2, 2, 3]),
                             array_ops.ones([1, 2]))


def get_sync_optimizer():
  return sync_replicas_optimizer.SyncReplicasOptimizer(
      gradient_descent.GradientDescentOptimizer(learning_rate=1.0),
      replicas_to_aggregate=1)


class GANModelTest(test.TestCase, parameterized.TestCase):
  """Tests for `gan_model`."""

  @parameterized.named_parameters(
      ('gan', get_gan_model, namedtuples.GANModel),
      ('callable_gan', get_callable_gan_model, namedtuples.GANModel),
      ('infogan', get_infogan_model, namedtuples.InfoGANModel),
      ('callable_infogan', get_callable_infogan_model,
       namedtuples.InfoGANModel),
      ('acgan', get_acgan_model, namedtuples.ACGANModel),
      ('callable_acgan', get_callable_acgan_model, namedtuples.ACGANModel),
      ('cyclegan', get_cyclegan_model, namedtuples.CycleGANModel),
      ('callable_cyclegan', get_callable_cyclegan_model,
       namedtuples.CycleGANModel),
      ('stargan', get_stargan_model, namedtuples.StarGANModel),
      ('callabel_stargan', get_callable_stargan_model, namedtuples.StarGANModel)
  )
  def test_output_type(self, create_fn, expected_tuple_type):
    """Test that output type is as expected."""
    self.assertIsInstance(create_fn(), expected_tuple_type)

  def test_no_shape_check(self):

    def dummy_generator_model(_):
      return (None, None)

    def dummy_discriminator_model(data, conditioning):  # pylint: disable=unused-argument
      return 1

    with self.assertRaisesRegexp(AttributeError, 'object has no attribute'):
      train.gan_model(
          dummy_generator_model,
          dummy_discriminator_model,
          real_data=array_ops.zeros([1, 2]),
          generator_inputs=array_ops.zeros([1]),
          check_shapes=True)
    train.gan_model(
        dummy_generator_model,
        dummy_discriminator_model,
        real_data=array_ops.zeros([1, 2]),
        generator_inputs=array_ops.zeros([1]),
        check_shapes=False)


class StarGANModelTest(test.TestCase):
  """Tests for `stargan_model`."""

  @staticmethod
  def create_input_and_label_tensor(batch_size, img_size, c_size, num_domains):
    input_tensor_list = []
    label_tensor_list = []
    for _ in range(num_domains):
      input_tensor_list.append(
          random_ops.random_uniform((batch_size, img_size, img_size, c_size)))
      domain_idx = random_ops.random_uniform(
          [batch_size], minval=0, maxval=num_domains, dtype=dtypes.int32)
      label_tensor_list.append(array_ops.one_hot(domain_idx, num_domains))
    return input_tensor_list, label_tensor_list

  def test_generate_stargan_random_domain_target(self):
    batch_size = 8
    domain_numbers = 3

    target_tensor = train._generate_stargan_random_domain_target(
        batch_size, domain_numbers)

    with self.test_session() as sess:
      targets = sess.run(target_tensor)
      self.assertTupleEqual((batch_size, domain_numbers), targets.shape)
      for target in targets:
        self.assertEqual(1, np.sum(target))
        self.assertEqual(1, np.max(target))

  def test_stargan_model_output_type(self):
    batch_size = 2
    img_size = 16
    c_size = 3
    num_domains = 5

    input_tensor, label_tensor = StarGANModelTest.create_input_and_label_tensor(
        batch_size, img_size, c_size, num_domains)
    model = train.stargan_model(
        generator_fn=stargan_generator_model,
        discriminator_fn=stargan_discriminator_model,
        input_data=input_tensor,
        input_data_domain_label=label_tensor)

    self.assertIsInstance(model, namedtuples.StarGANModel)
    self.assertTrue(isinstance(model.discriminator_variables, list))
    self.assertTrue(isinstance(model.generator_variables, list))
    self.assertIsInstance(model.discriminator_scope,
                          variable_scope.VariableScope)
    self.assertTrue(model.generator_scope, variable_scope.VariableScope)
    self.assertTrue(callable(model.discriminator_fn))
    self.assertTrue(callable(model.generator_fn))

  def test_stargan_model_generator_output(self):
    batch_size = 2
    img_size = 16
    c_size = 3
    num_domains = 5

    input_tensor, label_tensor = StarGANModelTest.create_input_and_label_tensor(
        batch_size, img_size, c_size, num_domains)
    model = train.stargan_model(
        generator_fn=stargan_generator_model,
        discriminator_fn=stargan_discriminator_model,
        input_data=input_tensor,
        input_data_domain_label=label_tensor)

    with self.test_session(use_gpu=True) as sess:

      sess.run(variables.global_variables_initializer())

      input_data, generated_data, reconstructed_data = sess.run(
          [model.input_data, model.generated_data, model.reconstructed_data])
      self.assertTupleEqual(
          (batch_size * num_domains, img_size, img_size, c_size),
          input_data.shape)
      self.assertTupleEqual(
          (batch_size * num_domains, img_size, img_size, c_size),
          generated_data.shape)
      self.assertTupleEqual(
          (batch_size * num_domains, img_size, img_size, c_size),
          reconstructed_data.shape)

  def test_stargan_model_discriminator_output(self):
    batch_size = 2
    img_size = 16
    c_size = 3
    num_domains = 5

    input_tensor, label_tensor = StarGANModelTest.create_input_and_label_tensor(
        batch_size, img_size, c_size, num_domains)
    model = train.stargan_model(
        generator_fn=stargan_generator_model,
        discriminator_fn=stargan_discriminator_model,
        input_data=input_tensor,
        input_data_domain_label=label_tensor)

    with self.test_session(use_gpu=True) as sess:

      sess.run(variables.global_variables_initializer())

      disc_input_data_source_pred, disc_gen_data_source_pred = sess.run([
          model.discriminator_input_data_source_predication,
          model.discriminator_generated_data_source_predication
      ])
      self.assertEqual(1, len(disc_input_data_source_pred.shape))
      self.assertEqual(batch_size * num_domains,
                       disc_input_data_source_pred.shape[0])
      self.assertEqual(1, len(disc_gen_data_source_pred.shape))
      self.assertEqual(batch_size * num_domains,
                       disc_gen_data_source_pred.shape[0])

      input_label, disc_input_label, gen_label, disc_gen_label = sess.run([
          model.input_data_domain_label,
          model.discriminator_input_data_domain_predication,
          model.generated_data_domain_target,
          model.discriminator_generated_data_domain_predication
      ])
      self.assertTupleEqual((batch_size * num_domains, num_domains),
                            input_label.shape)
      self.assertTupleEqual((batch_size * num_domains, num_domains),
                            disc_input_label.shape)
      self.assertTupleEqual((batch_size * num_domains, num_domains),
                            gen_label.shape)
      self.assertTupleEqual((batch_size * num_domains, num_domains),
                            disc_gen_label.shape)


class GANLossTest(test.TestCase, parameterized.TestCase):
  """Tests for `gan_loss`."""

  @parameterized.named_parameters(
      ('gan', get_gan_model),
      ('callable_gan', get_callable_gan_model),
      ('infogan', get_infogan_model),
      ('callable_infogan', get_callable_infogan_model),
      ('acgan', get_acgan_model),
      ('callable_acgan', get_callable_acgan_model),