# Copyright 2022 DeepMind Technologies Limited
#
# 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.

"""A test dataset that may be constructed entirely in-memory.

The test dataset consists of images. The task is to predict the dominant color
(red, green, or blue).
"""

import functools
from typing import Iterator, Optional, Tuple

from dm_nevis.benchmarker.datasets import dataset_builders
from dm_nevis.benchmarker.datasets import datasets
from dm_nevis.benchmarker.datasets import tasks
import numpy as np
import tensorflow as tf


DATASET_NAME = 'test_dataset'
DEFAULT_IMAGE_SIZE = 32
DEFAULT_SHUFFLE_BUFFER_SIZE = 10
SPLIT_DEFAULTS = {
    'test': {
        'noise_level': 64,
        'num_examples': 100,
        'seed': 0,
    },
    'train': {
        'noise_level': 32,
        'num_examples': 1000,
        'seed': 1,
    },
    'val': {
        'noise_level': 32,
        'num_examples': 100,
        'seed': 2,
    }
}


def get_dataset(
    split: str,
    *,
    start: Optional[int] = None,
    end: Optional[int] = None,
    image_size: int = DEFAULT_IMAGE_SIZE,
    task_kind: tasks.TaskKind = tasks.TaskKind.CLASSIFICATION
) -> datasets.Dataset:
  """Gets the test dataset.

  The task is to classify the color of a flood filled image as Red, Green,
  or Blue. The images in the dataset have been flood filled to one of these
  colors, and then perturbed with uniform noise.

  Args:
    split: One of ["test", "train", "val"].
    start: An optional offset from the beginning of the dataset to start at.
    end: An optional offset from the beggning of the dataset to end at.
    image_size: Select the size of the image to be returned.
    task_kind: The task kind to use.

  Returns:
    A dataset containing at most end - start images.
  """

  kwargs = SPLIT_DEFAULTS[split]
  builder_fn = _make_builder_fn(
      outer_start=start,
      outer_end=end,
      image_size=image_size,
      task_kind=task_kind,
      **kwargs,
  )

  return datasets.Dataset(
      task_key=_task_key(task_kind),
      builder_fn=builder_fn,
      num_examples=_compute_num_examples(kwargs['num_examples'], start, end))


def _make_builder_fn(
    outer_start: Optional[int],
    outer_end: Optional[int],
    image_size: int,
    seed: int,
    num_examples: int,
    noise_level: int,
    task_kind: tasks.TaskKind,
) -> datasets.DatasetBuilderFn:
  """Constructs a builder function for the test dataset."""

  def builder_fn(*,
                 shuffle: bool,
                 start: Optional[int] = None,
                 end: Optional[int] = None) -> tf.data.Dataset:

    start, end = dataset_builders.combine_indices(outer_start, start, outer_end,
                                                  end)

    def gen():
      yield from _generate_data(
          seed,
          num_examples,
          image_size,
          noise_level,
      )

    ds = tf.data.Dataset.from_generator(
        gen,
        output_signature=(tf.TensorSpec(
            shape=(image_size, image_size, 3),
            dtype=tf.uint8), tf.TensorSpec(shape=(), dtype=tf.int32)))

    if start is not None:
      ds = ds.skip(start)
    if end is not None:
      ds = ds.take(end - (start or 0))
    if shuffle:
      ds = ds.shuffle(buffer_size=DEFAULT_SHUFFLE_BUFFER_SIZE)

    ds = ds.map(functools.partial(_to_minibatch, task_kind=task_kind))

    return ds

  return builder_fn


def _generate_data(seed: int, num_examples: int, image_size: int,
                   noise_level: int) -> Iterator[Tuple[tf.Tensor, tf.Tensor]]:
  """Generates data for the dataset.

  Args:
    seed: A seed ensures that each run produces deterministic output.
    num_examples: The number of examples that the generator should produce.
    image_size: The image size of the resulting data.
    noise_level: Noise added to the images, to make the task a little more
      challenging.

  Yields:
    A generator over (image, label) tuples. The image is of shape
    (image_size, image_size, 3) and in the range [0, 255].
    Each image is generated by flood filling a single channel (R, G, or B)
    and the associated label is the channel that was flood filled.
    Noise in the range (-noise_level, noise_level) is added to the image,
    to make the classification task a little more challenging.
    The label is a value from {0, 1, 2} representing the dominant color
    of the returned image.
  """

  gen = np.random.default_rng(seed=seed)

  for _ in range(num_examples):
    color = gen.integers(0, 3)
    img = np.zeros((image_size, image_size, 3))
    img[:, :, color] = 255
    if noise_level > 0:
      img += gen.integers(
          -noise_level, noise_level, size=(image_size, image_size, 3))
    img = img.clip(0, 255)
    yield tf.constant(img, dtype=tf.uint8), color


def _to_minibatch(
    image: tf.Tensor,
    label: tf.Tensor,
    task_kind: tasks.TaskKind,
) -> datasets.MiniBatch:
  """Create a minibatch from the generated example."""

  if task_kind is tasks.TaskKind.CLASSIFICATION:
    return datasets.MiniBatch(
        image=image,
        label=label,
        multi_label_one_hot=None,
    )

  elif task_kind is tasks.TaskKind.MULTI_LABEL_CLASSIFICATION:
    label_one_hot = tf.one_hot([label], depth=3)[0]
    return datasets.MiniBatch(
        image=image,
        label=None,
        multi_label_one_hot=label_one_hot,
    )

  else:
    raise ValueError(f'Unsupported task kind: {task_kind}')


def _compute_num_examples(n: int, start: Optional[int],
                          end: Optional[int]) -> int:
  """Computes number of examples from known length and optional offsets."""
  start = start or 0
  end = end or n
  return min(n, end - start)


def _task_key(task_kind: tasks.TaskKind) -> tasks.TaskKey:
  """Creates a task key given the task kind."""

  if task_kind is tasks.TaskKind.MULTI_LABEL_CLASSIFICATION:
    return tasks.TaskKey(
        DATASET_NAME,
        tasks.TaskKind.MULTI_LABEL_CLASSIFICATION,
        tasks.MultiLabelClassificationMetadata(num_classes=3),
    )
  elif task_kind is tasks.TaskKind.CLASSIFICATION:
    return tasks.TaskKey(
        DATASET_NAME,
        tasks.TaskKind.CLASSIFICATION,
        tasks.ClassificationMetadata(num_classes=3),
    )
  else:
    raise ValueError(f'Unsupported task kind: {task_kind}')
