# coding=utf-8
# Copyright 2023 The TensorFlow GAN Authors.
#
# 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.

"""Definitions of generator functions."""

import tensorflow.compat.v1 as tf
import tensorflow_gan as tfgan
from tensorflow_gan.examples.self_attention_estimator import ops


def make_z_normal(num_batches, batch_size, z_dim):
  """Make random noise tensors with normal distribution.

  Args:
    num_batches: copies of batches
    batch_size: the batch_size for z
    z_dim: The dimension of the z (noise) vector.
  Returns:
    zs:  noise tensors.
  """
  shape = [num_batches, batch_size, z_dim]
  z = tf.random.normal(shape, name='z0', dtype=tf.float32)
  return z


def make_class_labels(batch_size, num_classes):
  """Generate class labels for generation."""
  # Uniform distribution.
  # TODO(augustusodena) Use true distribution of ImageNet classses.
  gen_class_logits = tf.zeros((batch_size, num_classes))
  gen_class_ints = tf.random.categorical(logits=gen_class_logits, num_samples=1)
  gen_class_ints.shape.assert_has_rank(2)
  gen_class_ints = tf.squeeze(gen_class_ints, axis=1)
  gen_class_ints.shape.assert_has_rank(1)

  return gen_class_ints


def usample(x):
  """Upsamples the input volume.

  Args:
    x: The 4D input tensor.
  Returns:
    An upsampled version of the input tensor.
  """
  # Allow the batch dimension to be unknown at graph build time.
  _, image_height, image_width, n_channels = x.shape.as_list()
  # Add extra degenerate dimension after the dimensions corresponding to the
  # rows and columns.
  expanded_x = tf.expand_dims(tf.expand_dims(x, axis=2), axis=4)
  # Duplicate data in the expanded dimensions.
  after_tile = tf.tile(expanded_x, [1, 1, 2, 1, 2, 1])
  return tf.reshape(after_tile,
                    [-1, image_height * 2, image_width * 2, n_channels])


def block(x, labels, out_channels, num_classes, name, training=True):
  """Builds the residual blocks used in the generator.

  Args:
    x: The 4D input tensor.
    labels: The labels of the class we seek to sample from.
    out_channels: Integer number of features in the output layer.
    num_classes: Integer number of classes in the labels.
    name: The variable scope name for the block.
    training: Whether this block is for training or not.
  Returns:
    A `Tensor` representing the output of the operation.
  """
  with tf.variable_scope(name):
    labels_onehot = tf.one_hot(labels, num_classes)
    x_0 = x
    x = tf.nn.relu(tfgan.tpu.batch_norm(x, training, labels_onehot,
                                        name='cbn_0'))
    x = usample(x)
    x = ops.snconv2d(x, out_channels, 3, 3, 1, 1, training, 'snconv1')
    x = tf.nn.relu(tfgan.tpu.batch_norm(x, training, labels_onehot,
                                        name='cbn_1'))
    x = ops.snconv2d(x, out_channels, 3, 3, 1, 1, training, 'snconv2')

    x_0 = usample(x_0)
    x_0 = ops.snconv2d(x_0, out_channels, 1, 1, 1, 1, training, 'snconv3')

    return x_0 + x


def generator(zs, target_class, gf_dim, num_classes, training=True):
  """Builds the generator segment of the graph, going from z -> G(z).

  Args:
    zs: Tensor representing the latent variables.
    target_class: The class from which we seek to sample.
    gf_dim: The gf dimension.
    num_classes: Number of classes in the labels.
    training: Whether in train mode or not. This affects things like batch
      normalization and spectral normalization.

  Returns:
    - The output layer of the generator.
    - A list containing all trainable varaibles defined by the model.
  """
  with tf.variable_scope('generator', reuse=tf.AUTO_REUSE) as gen_scope:

    act0 = ops.snlinear(
        zs, gf_dim * 16 * 4 * 4, training=training, name='g_snh0')
    act0 = tf.reshape(act0, [-1, 4, 4, gf_dim * 16])

    # pylint: disable=line-too-long
    act1 = block(act0, target_class, gf_dim * 16, num_classes, 'g_block1', training)  # 8
    act2 = block(act1, target_class, gf_dim * 8, num_classes, 'g_block2', training)  # 16
    act3 = block(act2, target_class, gf_dim * 4, num_classes, 'g_block3', training)  # 32
    act3 = ops.sn_non_local_block_sim(act3, training, name='g_ops')  # 32
    act4 = block(act3, target_class, gf_dim * 2, num_classes, 'g_block4', training)  # 64
    act5 = block(act4, target_class, gf_dim, num_classes, 'g_block5', training)  # 128
    act5 = tf.nn.relu(tfgan.tpu.batch_norm(act5, training, conditional_class_labels=None, name='g_bn'))
    act6 = ops.snconv2d(act5, 3, 3, 3, 1, 1, training, 'g_snconv_last')
    out = tf.nn.tanh(act6)
  var_list = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, gen_scope.name)
  return out, var_list
