# 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 and preprocessing the ImageNet data."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

from absl import flags
import tensorflow.compat.v1 as tf  # tf
import tensorflow_datasets as tfds

from tensorflow_gan.examples import compat_utils

IMG_SIZE = 128

flags.DEFINE_string('imagenet_data_dir', None,
                    'A directory for TFDS ImageNet. If `None`, use default.')


def provide_dataset(batch_size, shuffle_buffer_size, split='train'):
  """Provides dataset of ImageNet digits that were preprocessed by the Red Team.

  Args:
    batch_size: The number of images in each batch.
    shuffle_buffer_size: The number of records to load before shuffling. Larger
      means more likely randomization.
    split: A tfds split. If 'train', dataset is shuffled. Otherwise, it's
      deterministic.
  Returns:
    A dataset of num_batches batches of size batch_size of images and labels.
  """
  shuffle = (split in ['train', tfds.Split.TRAIN])
  dataset = _load_imagenet_dataset(split, flags.FLAGS.imagenet_data_dir,
                                   shuffle_files=shuffle)
  if shuffle:
    dataset = dataset.apply(
        tf.data.experimental.shuffle_and_repeat(shuffle_buffer_size))
  else:
    dataset = dataset.repeat()
  dataset = (dataset.map(_preprocess_dataset_record_fn(IMG_SIZE),
                         num_parallel_calls=16 if shuffle else None)
             .batch(batch_size, drop_remainder=True))
  dataset = dataset.prefetch(tf.data.experimental.AUTOTUNE)
  return dataset


def provide_data(batch_size,
                 num_batches,
                 shuffle_buffer_size,
                 split='train'):
  """Provides batches of ImageNet digits that were preprocessed by the Red Team.

  Args:
    batch_size: The number of images in each batch.
    num_batches: The number of batches to return.
    shuffle_buffer_size: The number of records to load before shuffling. Larger
      means more likely randomization.
    split: A tfds split.

  Returns:
    A list of num_batches batches of size batch_size. Each element in the
    returned list is a tuple of a batch of images and a batch of the respective
    labels.
  """
  dataset = provide_dataset(batch_size * num_batches, shuffle_buffer_size,
                            split)
  iterator = tf.data.make_one_shot_iterator(dataset)
  images, labels = iterator.get_next()
  images = tf.reshape(
      images, shape=[num_batches, batch_size, IMG_SIZE, IMG_SIZE, 3])
  labels = tf.reshape(labels, shape=[num_batches, batch_size, 1])
  batches = list(
      zip(
          tf.unstack(images, num=num_batches),
          tf.unstack(labels, num=num_batches)))
  return batches


def _load_imagenet_dataset(split, data_dir=None, shuffle_files=False):
  return tfds.load('imagenet2012:5.*.*', split=split, data_dir=data_dir,
                   shuffle_files=shuffle_files)


def _preprocess_dataset_record_fn(image_size):
  """Returns function for processing the elements of the imagenet dataset."""

  def _process_record(record):
    """Takes the largest central square and resamples to image_size."""
    # Based on
    # https://github.com/openai/improved-gan/blob/master/imagenet/convert_imagenet_to_records.py
    image = record['image']
    image_shape = tf.cast(tf.shape(input=image), tf.float32)
    box_size = tf.math.minimum(image_shape[0], image_shape[1])
    # Since we assume the box is centered we have:
    # 2 * box_x_min + box_size == box_width,
    # 2 * box_y_min + box_size == box_height.
    # tf.math.ceil is used for consistency with the improved-gan implementation.
    box_y_min = tf.math.ceil(0.5 * (image_shape[0] - box_size))
    box_x_min = tf.math.ceil(0.5 * (image_shape[1] - box_size))
    box_y_max = box_y_min + box_size - 1
    box_x_max = box_x_min + box_size - 1
    # Normalize with the inverse of the trasform done by crop_and_resize.
    normalized_y_min = box_y_min / (image_shape[0] - 1)
    normalized_x_min = box_x_min / (image_shape[1] - 1)
    normalized_y_max = box_y_max / (image_shape[0] - 1)
    normalized_x_max = box_x_max / (image_shape[1] - 1)
    image = compat_utils.crop_and_resize([image],
                                         boxes=[[
                                             normalized_y_min, normalized_x_min,
                                             normalized_y_max, normalized_x_max
                                         ]],
                                         box_ind=[0],
                                         crop_size=[image_size, image_size])
    # crop_and_resize returns a tensor of type tf.float32.
    image = tf.squeeze(image, axis=0)
    image = image * (2. / 255) - 1.
    label = tf.cast(record['label'], tf.int32)
    return image, label

  return _process_record
