# coding=utf-8
# Copyright 2022 The Google Research 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.

# Lint as: python3
"""Helper functions to pre-process data inputs and create data iterator.

Predictions must be stored in the same sequence for each checkpoint.
This requires a tailored data input script to preserve ordering.
"""

import numpy as np
import tensorflow.compat.v1 as tf
from pruning_identified_exemplars.utils import preprocessing_helper


def input_fn(params):
  """Input function from cloud tpu resnet model pre-processing."""

  def parser_tfrecords(serialized_example):
    """Parses a single tf.Example into image and label tensors."""
    data = {}
    if params['test_small_sample']:
      data['image_raw'] = serialized_example
      data['label'] = tf.constant(0, tf.int32)
      data['human_label'] = tf.constant('human_label', tf.string)
      data['key_'] = tf.constant('key', tf.string)
      return data
    else:
      features = tf.parse_single_example(
          serialized_example,
          features={
              'key_':
                  tf.FixedLenFeature([], dtype=tf.string, default_value=''),
              'image/encoded':
                  tf.FixedLenFeature((), tf.string, default_value=''),
              'image/class/label':
                  tf.FixedLenFeature([1], dtype=tf.int64, default_value=-1),
              'image/class/text':
                  tf.FixedLenFeature([], dtype=tf.string, default_value=''),
          })

      image_bytes = tf.reshape(features['image/encoded'], shape=[])
      label = tf.cast(
          tf.reshape(features['image/class/label'], shape=[]),
          dtype=tf.int32) - 1
      image = tf.image.decode_jpeg(features['image/encoded'], 3)

      human_label = tf.cast(
          tf.reshape(features['image/class/text'], shape=[]), dtype=tf.string)
      if params['task'] == 'imagenet_training':
        # training is set to false in prediction mode
        image = preprocessing_helper.preprocess_image(
            image=image, image_size=224, is_training=True)
      else:
        # training is set to false in prediction mode
        image = preprocessing_helper.preprocess_image(
            image=image, image_size=224, is_training=False)

      if params['task'] == 'pie_dataset_gen':
        data['image_raw'] = image_bytes
      else:
        data['image_raw'] = image
      data['label'] = label
      data['human_label'] = human_label
      data['key_'] = tf.reshape(features['key_'], shape=[])
      return data

  def _get_null_input(data):
    """Returns a null image (all black pixels)."""
    del data
    return tf.zeros([224, 224, 3], tf.float32)

  data_dir = params['data_dir']
  # batch_size is set to 1 to avoid any images being dropped
  if params['task'] in ['imagenet_predictions', 'pie_dataset_gen']:
    batch_size = 1
    sloppy_shuffle = False
  else:
    batch_size = params['batch_size']
    sloppy_shuffle = params['sloppy_shuffle']

  if params['test_small_sample']:
    dataset = tf.data.Dataset.range(1).repeat().map(_get_null_input)
  else:
    dataset = tf.data.Dataset.list_files(data_dir, shuffle=False)

    def fetch_dataset(filename):
      dataset = tf.data.TFRecordDataset(filename)
      return dataset

    dataset = dataset.apply(
        tf.data.experimental.parallel_interleave(
            fetch_dataset,
            cycle_length=params['num_cores'],
            sloppy=sloppy_shuffle))

  if params['mode'] == 'train':
    dataset = dataset.shuffle(buffer_size=1024)
    dataset = dataset.repeat()

  dataset = dataset.map(parser_tfrecords, num_parallel_calls=64)
  dataset = dataset.prefetch(batch_size)
  dataset = dataset.batch(batch_size, drop_remainder=sloppy_shuffle)
  dataset = dataset.make_one_shot_iterator().get_next()

  if params['task'] == 'imagenet_predictions':
    return ([
        tf.reshape(dataset['image_raw'], [batch_size, 224, 224, 3]),
        tf.reshape(dataset['label'], [batch_size])
    ], tf.reshape(dataset['label'], [batch_size]))
  else:
    return dataset


def image_to_tfexample(key,
                       raw_image,
                       label,
                       stored_class_label,
                       image_index=None,
                       baseline_n=None,
                       variant_n=None,
                       pruning_fraction=None,
                       baseline_mean_probability=None,
                       variant_mean_probability=None,
                       variant_mode_label=None,
                       baseline_mode_label=None,
                       test_small_sample=False):
  """Generates a serialized TF-example."""

  if test_small_sample:
    type_raw_image = tf.train.Feature(
        float_list=tf.train.FloatList(
            value=np.reshape(raw_image.astype(np.float32), [-1])))
  else:
    type_raw_image = tf.train.Feature(
        bytes_list=tf.train.BytesList(value=[raw_image]))
  return tf.train.Example(
      features=tf.train.Features(
          feature={
              'key':
                  tf.train.Feature(bytes_list=tf.train.BytesList(value=[key])),
              'image/encoded':
                  type_raw_image,
              'image/class/label':
                  tf.train.Feature(
                      int64_list=tf.train.Int64List(value=[label])),
              'stored_class_label':
                  tf.train.Feature(
                      int64_list=tf.train.Int64List(
                          value=[stored_class_label])),
              'image_index':
                  tf.train.Feature(
                      int64_list=tf.train.Int64List(value=[image_index])),
              'baseline_n':
                  tf.train.Feature(
                      int64_list=tf.train.Int64List(value=[baseline_n])),
              'variant_n':
                  tf.train.Feature(
                      int64_list=tf.train.Int64List(value=[variant_n])),
              'pruning_fraction':
                  tf.train.Feature(
                      float_list=tf.train.FloatList(value=[pruning_fraction])),
              'baseline_mean_probability':
                  tf.train.Feature(
                      float_list=tf.train.FloatList(
                          value=[baseline_mean_probability])),
              'variant_mean_probability':
                  tf.train.Feature(
                      float_list=tf.train.FloatList(
                          value=[variant_mean_probability])),
              'variant_mode_label':
                  tf.train.Feature(
                      int64_list=tf.train.Int64List(
                          value=[variant_mode_label])),
              'baseline_mode_label':
                  tf.train.Feature(
                      int64_list=tf.train.Int64List(
                          value=[baseline_mode_label])),
          }))
