# coding=utf-8
# Copyright 2020 The TensorFlow Datasets 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
"""Corrupted MNIST Dataset.

MNISTCorrupted is a dataset generated by adding 15 corruptions to the test
images in the MNIST dataset. This dataset wraps the static, corrupted MNIST
test images uploaded by the original authors.
"""

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

import os

import numpy as np
import tensorflow.compat.v2 as tf
from tensorflow_datasets.image_classification import mnist
import tensorflow_datasets.public_api as tfds

_DESCRIPTION = """\
MNISTCorrupted is a dataset generated by adding 15 corruptions to the test
images in the MNIST dataset. This dataset wraps the static, corrupted MNIST
test images uploaded by the original authors
"""

_CITATION = """
@article{mu2019mnist,
  title={MNIST-C: A Robustness Benchmark for Computer Vision},
  author={Mu, Norman and Gilmer, Justin},
  journal={arXiv preprint arXiv:1906.02337},
  year={2019}
}
"""

_DOWNLOAD_URL = 'https://zenodo.org/record/3239543/files/mnist_c.zip'
_CORRUPTIONS = [
    'identity',
    'shot_noise',
    'impulse_noise',
    'glass_blur',
    'motion_blur',
    'shear',
    'scale',
    'rotate',
    'brightness',
    'translate',
    'stripe',
    'fog',
    'spatter',
    'dotted_line',
    'zigzag',
    'canny_edges',
]
_DIRNAME = 'mnist_c'
_TRAIN_IMAGES_FILENAME = 'train_images.npy'
_TEST_IMAGES_FILENAME = 'test_images.npy'
_TRAIN_LABELS_FILENAME = 'train_labels.npy'
_TEST_LABELS_FILENAME = 'test_labels.npy'


class MNISTCorruptedConfig(tfds.core.BuilderConfig):
  """BuilderConfig for MNISTcorrupted."""

  @tfds.core.disallow_positional_args
  def __init__(self, corruption_type, **kwargs):
    """Constructor.

    Args:
      corruption_type: string, name of corruption from _CORRUPTIONS.
      **kwargs: keyword arguments forwarded to super.
    """
    super(MNISTCorruptedConfig, self).__init__(**kwargs)
    self.corruption = corruption_type


def _make_builder_configs():
  """Construct a list of BuilderConfigs.

  Construct a list of 15 MNISTCorruptedConfig objects, corresponding to
  the 15 corruption types.

  Returns:
    A list of 15 MNISTCorruptedConfig objects.
  """
  config_list = []
  for corruption in _CORRUPTIONS:
    config_list.append(
        MNISTCorruptedConfig(
            name=corruption,
            version=tfds.core.Version(
                '1.0.0',
                'New split API (https://tensorflow.org/datasets/splits)'),
            description='Corruption method: ' + corruption,
            corruption_type=corruption,
        ))
  return config_list


class MNISTCorrupted(tfds.core.GeneratorBasedBuilder):
  """Corrupted MNIST dataset."""
  BUILDER_CONFIGS = _make_builder_configs()

  def _info(self):
    """Returns basic information of dataset.

    Returns:
      tfds.core.DatasetInfo.
    """
    return tfds.core.DatasetInfo(
        builder=self,
        description=_DESCRIPTION,
        features=tfds.features.FeaturesDict({
            'image':
                tfds.features.Image(shape=mnist.MNIST_IMAGE_SHAPE),
            'label':
                tfds.features.ClassLabel(num_classes=mnist.MNIST_NUM_CLASSES),
        }),
        supervised_keys=('image', 'label'),
        homepage='https://github.com/google-research/mnist-c',
        citation=_CITATION)

  def _split_generators(self, dl_manager):
    """Return the train, test split of MNIST-C.

    Args:
      dl_manager: download manager object.

    Returns:
      train split, test split.
    """
    path = dl_manager.download_and_extract(_DOWNLOAD_URL)
    return [
        tfds.core.SplitGenerator(
            name=tfds.Split.TRAIN,
            gen_kwargs={
                'data_dir': os.path.join(path, _DIRNAME),
                'is_train': True
            }),
        tfds.core.SplitGenerator(
            name=tfds.Split.TEST,
            gen_kwargs={
                'data_dir': os.path.join(path, _DIRNAME),
                'is_train': False
            }),
    ]

  def _generate_examples(self, data_dir, is_train):
    """Generate corrupted MNIST data.

    Apply corruptions to the raw images according to self.corruption_type.

    Args:
      data_dir: root directory of downloaded dataset
      is_train: whether to return train images or test images

    Yields:
      dictionary with image file and label.
    """
    corruption = self.builder_config.corruption

    if is_train:
      images_file = os.path.join(data_dir, corruption, _TRAIN_IMAGES_FILENAME)
      labels_file = os.path.join(data_dir, corruption, _TRAIN_LABELS_FILENAME)
    else:
      images_file = os.path.join(data_dir, corruption, _TEST_IMAGES_FILENAME)
      labels_file = os.path.join(data_dir, corruption, _TEST_LABELS_FILENAME)

    with tf.io.gfile.GFile(labels_file, mode='rb') as f:
      labels = np.load(f)

    with tf.io.gfile.GFile(images_file, mode='rb') as f:
      images = np.load(f)

    for i, (image, label) in enumerate(zip(images, labels)):
      yield i, {
          'image': image,
          'label': label,
      }