# Copyright 2019, Google LLC.
#
# 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.
"""Preprocessing library for EMNIST baseline tasks."""

from collections.abc import Callable
from typing import Optional

import tensorflow as tf

from tensorflow_federated.python.simulation.baselines import client_spec

MAX_CLIENT_DATASET_SIZE = 418


def _reshape_for_character_recognition(element):
  return tf.expand_dims(element['pixels'], axis=-1), element['label']


def _reshape_for_autoencoder(element):
  x = 1 - tf.reshape(element['pixels'], (-1, 28 * 28))
  return (x, x)


def create_preprocess_fn(
    preprocess_spec: client_spec.ClientSpec,
    emnist_task: str = 'character_recognition',
    num_parallel_calls: tf.Tensor = tf.data.experimental.AUTOTUNE,
    debug_seed: Optional[int] = None,
) -> Callable[[tf.data.Dataset], tf.data.Dataset]:
  """Creates a preprocessing function for EMNIST client datasets.

  The preprocessing shuffles, repeats, batches, and then reshapes, using
  the `shuffle`, `repeat`, `take`, `batch`, and `map` attributes of a
  `tf.data.Dataset`, in that order.

  Args:
    preprocess_spec: A `tff.simulation.baselines.ClientSpec` containing
      information on how to preprocess clients.
    emnist_task: A string indicating the EMNIST task being performed. Must be
      one of 'character_recognition' or 'autoencoder'. If the former, then
      elements are mapped to tuples of the form (pixels, label), if the latter
      then elements are mapped to tuples of the form (pixels, pixels).
    num_parallel_calls: An integer representing the number of parallel calls
      used when performing `tf.data.Dataset.map`.
    debug_seed: An optional integer seed for deterministic shuffling and
      mapping. Intended for unittesting.

  Returns:
    A callable taking as input a `tf.data.Dataset`, and returning a
    `tf.data.Dataset` formed by preprocessing according to the input arguments.
  """
  shuffle_buffer_size = preprocess_spec.shuffle_buffer_size
  if shuffle_buffer_size is None:
    shuffle_buffer_size = MAX_CLIENT_DATASET_SIZE

  if emnist_task == 'character_recognition':
    mapping_fn = _reshape_for_character_recognition
  elif emnist_task == 'autoencoder':
    mapping_fn = _reshape_for_autoencoder
  else:
    raise ValueError(
        'emnist_task must be one of "character_recognition" or "autoencoder".'
    )

  def preprocess_fn(dataset: tf.data.Dataset) -> tf.data.Dataset:
    if shuffle_buffer_size > 1:
      dataset = dataset.shuffle(shuffle_buffer_size, seed=debug_seed)
    if preprocess_spec.num_epochs > 1:
      dataset = dataset.repeat(preprocess_spec.num_epochs)
    if preprocess_spec.max_elements is not None:
      dataset = dataset.take(preprocess_spec.max_elements)
    dataset = dataset.batch(preprocess_spec.batch_size, drop_remainder=False)
    return dataset.map(
        mapping_fn,
        num_parallel_calls=num_parallel_calls,
        deterministic=debug_seed is not None,
    )

  return preprocess_fn
