# 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.

"""Loss functions for training ESRGAN model.

The ESRGAN model makes use of three loss functions: pixel loss, perceptual
loss(vgg_loss) and adversarial loss. Adversarial loss for the model is
calculated using relativistic average loss which is pre-defined in TF-GAN.
"""

import tensorflow as tf


def pixel_loss(y_true, y_pred):
  """To calculate the L1 Loss.

  Args:
      y_true : Original HR image present in the dataset.
      y_pred : Image generated by the generator network.

  Returns:
      Mean absolute difference between the two images.
  """
  y_true = tf.cast(y_true, tf.float32)
  y_pred = tf.cast(y_pred, tf.float32)
  return tf.reduce_mean(tf.reduce_mean(tf.abs(y_true - y_pred), axis=0))


def vgg_loss(weight=None, input_shape=None):
  """Perceptual Loss calculation using pre-trained VGG19.

  Args:
      weight: Weights to be loaded.
      input_shape: Shape of input image.

  Returns:
      The calculated perceptual loss.
  """
  vgg_model = tf.keras.applications.vgg19.VGG19(
      input_shape=input_shape, weights=weight, include_top=False)

  for layer in vgg_model.layers:
    layer.trainable = False
  # Features before activation are used as mentioned in the paper.
  vgg_model.get_layer('block5_conv4').activation = lambda x: x
  vgg = tf.keras.Model(
      inputs=[vgg_model.input],
      outputs=[vgg_model.get_layer('block5_conv4').output])

  def loss(y_true, y_pred):
    return tf.compat.v1.losses.absolute_difference(vgg(y_true), vgg(y_pred))

  return loss
