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

"""Contains code for loading DIV2K dataset."""

import tensorflow as tf
import tensorflow_datasets as tfds


def random_flip(lr_img, hr_img):
  """Randomly flips LR and HR images for data augmentation."""
  random = tf.random.uniform(shape=(), maxval=1)
  flip_lr = tf.image.flip_left_right(lr_img)
  flip_hr = tf.image.flip_left_right(hr_img)
  return tf.cond(
      random < 0.5, lambda: (lr_img, hr_img), lambda: (flip_lr, flip_hr))


def random_rotate(lr_img, hr_img):
  """Randomly rotates LR and HR images for data augmentation."""
  random = tf.random.uniform(shape=(), maxval=4, dtype=tf.int32)
  return tf.image.rot90(lr_img, random), tf.image.rot90(hr_img, random)


def get_div2k_data(hparams,
                   name='div2k/bicubic_x4',
                   mode='train',
                   shuffle=True,
                   repeat_count=None,
                   augment=True):
  """Downloads and loads DIV2K dataset.

  Args:
      hparams : A named tuple to store different parameters.
      name : Name of the dataset  to be loaded using tfds.
      mode : Either 'train' or 'valid'.
      shuffle : Whether to shuffle the images in the dataset.
      repeat_count : Repetition of data during training.
      augment : Whether to augment the dataset or not.

  Returns:
      A tf.data.Dataset with pairs of LR image and HR image tensors.

  Raises:
        TypeError : If the data directory(data_dir) is not specified.
  """
  split = 'train' if mode == 'train' else 'validation'
  num_calls = tf.python.data.experimental.AUTOTUNE

  def scale(data):
    image = data['hr']
    hr_size = hparams.hr_dimension
    scale = hparams.scale

    hr_image = image
    hr_image = tf.image.resize(hr_image, [hr_size, hr_size])
    lr_image = tf.image.resize(
        hr_image, [hr_size // scale, hr_size // scale], method='bicubic')

    hr_image = tf.clip_by_value(hr_image, 0, 255)
    lr_image = tf.clip_by_value(lr_image, 0, 255)

    return lr_image, hr_image

  dataset = (
      tfds.load(name=name, split=split, data_dir=hparams.data_dir).map(
          scale, num_parallel_calls=num_calls).cache())

  if shuffle:
    dataset = dataset.shuffle(buffer_size=10000, reshuffle_each_iteration=True)

  if augment:
    dataset = dataset.map(random_flip)
    dataset = dataset.map(random_rotate)

  dataset = dataset.batch(hparams.batch_size)
  dataset = dataset.repeat(repeat_count)
  dataset = dataset.prefetch(buffer_size=num_calls)

  return dataset
