# coding=utf-8
# Copyright 2026 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.

"""MNIST, Fashion MNIST, KMNIST and EMNIST."""

import os

import numpy as np
from six.moves import urllib  # pyrefly: ignore[missing-source-for-stubs]
from tensorflow_datasets.core.utils.lazy_imports_utils import tensorflow as tf
import tensorflow_datasets.public_api as tfds

# MNIST constants
# CVDF mirror of http://yann.lecun.com/exdb/mnist/
_MNIST_URL = "https://storage.googleapis.com/cvdf-datasets/mnist/"
_MNIST_TRAIN_DATA_FILENAME = "train-images-idx3-ubyte.gz"
_MNIST_TRAIN_LABELS_FILENAME = "train-labels-idx1-ubyte.gz"
_MNIST_TEST_DATA_FILENAME = "t10k-images-idx3-ubyte.gz"
_MNIST_TEST_LABELS_FILENAME = "t10k-labels-idx1-ubyte.gz"
_MNIST_IMAGE_SIZE = 28
MNIST_IMAGE_SHAPE = (_MNIST_IMAGE_SIZE, _MNIST_IMAGE_SIZE, 1)
MNIST_NUM_CLASSES = 10
_TRAIN_EXAMPLES = 60000
_TEST_EXAMPLES = 10000

_MNIST_CITATION = """\
@article{lecun2010mnist,
  title={MNIST handwritten digit database},
  author={LeCun, Yann and Cortes, Corinna and Burges, CJ},
  journal={ATT Labs [Online]. Available: http://yann.lecun.com/exdb/mnist},
  volume={2},
  year={2010}
}
"""

_FASHION_MNIST_CITATION = """\
@article{DBLP:journals/corr/abs-1708-07747,
  author    = {Han Xiao and
               Kashif Rasul and
               Roland Vollgraf},
  title     = {Fashion-MNIST: a Novel Image Dataset for Benchmarking Machine Learning
               Algorithms},
  journal   = {CoRR},
  volume    = {abs/1708.07747},
  year      = {2017},
  url       = {http://arxiv.org/abs/1708.07747},
  archivePrefix = {arXiv},
  eprint    = {1708.07747},
  timestamp = {Mon, 13 Aug 2018 16:47:27 +0200},
  biburl    = {https://dblp.org/rec/bib/journals/corr/abs-1708-07747},
  bibsource = {dblp computer science bibliography, https://dblp.org}
}
"""

_K_MNIST_CITATION = """\
  @online{clanuwat2018deep,
  author       = {Tarin Clanuwat and Mikel Bober-Irizar and Asanobu Kitamoto and Alex Lamb and Kazuaki Yamamoto and David Ha},
  title        = {Deep Learning for Classical Japanese Literature},
  date         = {2018-12-03},
  year         = {2018},
  eprintclass  = {cs.CV},
  eprinttype   = {arXiv},
  eprint       = {cs.CV/1812.01718},
}
"""

_EMNIST_CITATION = """\
@article{cohen_afshar_tapson_schaik_2017,
    title={EMNIST: Extending MNIST to handwritten letters},
    DOI={10.1109/ijcnn.2017.7966217},
    journal={2017 International Joint Conference on Neural Networks (IJCNN)},
    author={Cohen, Gregory and Afshar, Saeed and Tapson, Jonathan and Schaik, Andre Van},
    year={2017}
}
"""


class MNIST(tfds.core.GeneratorBasedBuilder):
  """MNIST."""

  URL = _MNIST_URL

  VERSION = tfds.core.Version("3.0.1")

  def _info(self):
    return tfds.core.DatasetInfo(
        builder=self,
        description="The MNIST database of handwritten digits.",
        features=tfds.features.FeaturesDict({
            "image": tfds.features.Image(shape=MNIST_IMAGE_SHAPE),
            "label": tfds.features.ClassLabel(num_classes=MNIST_NUM_CLASSES),
        }),
        supervised_keys=("image", "label"),
        homepage="http://yann.lecun.com/exdb/mnist/",
        citation=_MNIST_CITATION,
    )

  def _split_generators(self, dl_manager):
    """Returns SplitGenerators."""
    # Download the full MNIST Database
    filenames = {
        "train_data": _MNIST_TRAIN_DATA_FILENAME,
        "train_labels": _MNIST_TRAIN_LABELS_FILENAME,
        "test_data": _MNIST_TEST_DATA_FILENAME,
        "test_labels": _MNIST_TEST_LABELS_FILENAME,
    }
    mnist_files = dl_manager.download_and_extract(
        {k: urllib.parse.urljoin(self.URL, v) for k, v in filenames.items()}
    )

    # MNIST provides TRAIN and TEST splits, not a VALIDATION split, so we only
    # write the TRAIN and TEST splits to disk.
    return [
        tfds.core.SplitGenerator(
            name=tfds.Split.TRAIN,  # pyrefly: ignore[missing-attribute]
            gen_kwargs=dict(
                num_examples=_TRAIN_EXAMPLES,
                data_path=mnist_files["train_data"],
                label_path=mnist_files["train_labels"],
            ),
        ),
        tfds.core.SplitGenerator(
            name=tfds.Split.TEST,  # pyrefly: ignore[missing-attribute]
            gen_kwargs=dict(
                num_examples=_TEST_EXAMPLES,
                data_path=mnist_files["test_data"],
                label_path=mnist_files["test_labels"],
            ),
        ),
    ]

  def _generate_examples(self, num_examples, data_path, label_path):
    """Generate MNIST examples as dicts.

    Args:
      num_examples (int): The number of example.
      data_path (str): Path to the data files
      label_path (str): Path to the labels

    Yields:
      Generator yielding the next examples
    """
    images = _extract_mnist_images(data_path, num_examples)
    labels = _extract_mnist_labels(label_path, num_examples)
    data = list(zip(images, labels))

    # Using index as key since data is always loaded in same order.
    for index, (image, label) in enumerate(data):
      record = {"image": image, "label": label}
      yield index, record


class FashionMNIST(MNIST):
  URL = "http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/"

  # TODO(afrozm): Try to inherit from MNIST's _info and mutate things as needed.
  def _info(self):
    return tfds.core.DatasetInfo(
        builder=self,
        description=(
            "Fashion-MNIST is a dataset of Zalando's article images "
            "consisting of a training set of 60,000 examples and a "
            "test set of 10,000 examples. Each example is a 28x28 "
            "grayscale image, associated with a label from 10 "
            "classes."
        ),
        features=tfds.features.FeaturesDict({
            "image": tfds.features.Image(shape=MNIST_IMAGE_SHAPE),
            "label": tfds.features.ClassLabel(
                names=[
                    "T-shirt/top",
                    "Trouser",
                    "Pullover",
                    "Dress",
                    "Coat",
                    "Sandal",
                    "Shirt",
                    "Sneaker",
                    "Bag",
                    "Ankle boot",
                ]
            ),
        }),
        supervised_keys=("image", "label"),
        homepage="https://github.com/zalandoresearch/fashion-mnist",
        citation=_FASHION_MNIST_CITATION,
    )


class KMNIST(MNIST):
  URL = "http://codh.rois.ac.jp/kmnist/dataset/kmnist/"

  def _info(self):
    return tfds.core.DatasetInfo(
        builder=self,
        description=(
            "Kuzushiji-MNIST is a drop-in replacement for the MNIST "
            "dataset (28x28 grayscale, 70,000 images), provided in "
            "the original MNIST format as well as a NumPy format. "
            "Since MNIST restricts us to 10 classes, we chose one "
            "character to represent each of the 10 rows of Hiragana "
            "when creating Kuzushiji-MNIST."
        ),
        features=tfds.features.FeaturesDict({
            "image": tfds.features.Image(shape=MNIST_IMAGE_SHAPE),
            "label": tfds.features.ClassLabel(
                names=[
                    "o",
                    "ki",
                    "su",
                    "tsu",
                    "na",
                    "ha",
                    "ma",
                    "ya",
                    "re",
                    "wo",
                ]
            ),
        }),
        supervised_keys=("image", "label"),
        homepage="http://codh.rois.ac.jp/kmnist/index.html.en",
        citation=_K_MNIST_CITATION,
    )


class EMNISTConfig(tfds.core.BuilderConfig):
  """BuilderConfig for EMNIST CONFIG."""

  def __init__(self, *, class_number, train_examples, test_examples, **kwargs):
    """BuilderConfig for EMNIST class number.

    Args:
      class_number: There are six different splits provided in this dataset. And
        have different class numbers.
      train_examples: number of train examples
      test_examples: number of test examples
      **kwargs: keyword arguments forwarded to super.
    """
    super(EMNISTConfig, self).__init__(**kwargs)
    self.class_number = class_number
    self.train_examples = train_examples
    self.test_examples = test_examples


class EMNIST(MNIST):
  """Emnist dataset."""

  URL = "https://biometrics.nist.gov/cs_links/EMNIST/gzip.zip"
  VERSION = tfds.core.Version("3.1.0")
  RELEASE_NOTES = {
      "3.0.0": "New split API (https://tensorflow.org/datasets/splits)",
      "3.1.0": "Updated broken download URL",
  }
  BUILDER_CONFIGS = [
      EMNISTConfig(
          name="byclass",
          class_number=62,
          train_examples=697932,
          test_examples=116323,
          description="EMNIST ByClass",
      ),
      EMNISTConfig(
          name="bymerge",
          class_number=47,
          train_examples=697932,
          test_examples=116323,
          description="EMNIST ByMerge",
      ),
      EMNISTConfig(
          name="balanced",
          class_number=47,
          train_examples=112800,
          test_examples=18800,
          description="EMNIST Balanced",
      ),
      EMNISTConfig(
          name="letters",
          class_number=37,
          train_examples=88800,
          test_examples=14800,
          description="EMNIST Letters",
      ),
      EMNISTConfig(
          name="digits",
          class_number=10,
          train_examples=240000,
          test_examples=40000,
          description="EMNIST Digits",
      ),
      EMNISTConfig(
          name="mnist",
          class_number=10,
          train_examples=60000,
          test_examples=10000,
          description="EMNIST MNIST",
      ),
  ]

  def _info(self):
    return tfds.core.DatasetInfo(
        builder=self,
        description=(
            "The EMNIST dataset is a set of handwritten character digits "
            "derived from the NIST Special Database 19 and converted to "
            "a 28x28 pixel image format and dataset structure that directly "
            "matches the MNIST dataset.\n\n"
            "Note: Like the original EMNIST data, images provided here are "
            "inverted horizontally and rotated 90 anti-clockwise. You can use "
            "`tf.transpose` within `ds.map` to convert the images to a "
            "human-friendlier format."
        ),
        features=tfds.features.FeaturesDict({
            "image": tfds.features.Image(shape=MNIST_IMAGE_SHAPE),
            "label": tfds.features.ClassLabel(
                num_classes=self.builder_config.class_number  # pyrefly: ignore[missing-attribute]
            ),
        }),
        supervised_keys=("image", "label"),
        homepage=(
            "https://www.nist.gov/itl/products-and-services/emnist-dataset"
        ),
        citation=_EMNIST_CITATION,
    )

  def _split_generators(self, dl_manager):
    filenames = {
        "train_data": "emnist-{}-train-images-idx3-ubyte.gz".format(
            self.builder_config.name  # pyrefly: ignore[missing-attribute]
        ),
        "train_labels": "emnist-{}-train-labels-idx1-ubyte.gz".format(
            self.builder_config.name  # pyrefly: ignore[missing-attribute]
        ),
        "test_data": "emnist-{}-test-images-idx3-ubyte.gz".format(
            self.builder_config.name  # pyrefly: ignore[missing-attribute]
        ),
        "test_labels": "emnist-{}-test-labels-idx1-ubyte.gz".format(
            self.builder_config.name  # pyrefly: ignore[missing-attribute]
        ),
    }

    dir_name = os.path.join(dl_manager.download_and_extract(self.URL), "gzip")
    extracted = dl_manager.extract(
        {k: os.path.join(dir_name, fname) for k, fname in filenames.items()}
    )

    return [
        tfds.core.SplitGenerator(
            name=tfds.Split.TRAIN,  # pyrefly: ignore[missing-attribute]
            gen_kwargs=dict(
                num_examples=self.builder_config.train_examples,  # pyrefly: ignore[missing-attribute]
                data_path=extracted["train_data"],
                label_path=extracted["train_labels"],
            ),
        ),
        tfds.core.SplitGenerator(
            name=tfds.Split.TEST,  # pyrefly: ignore[missing-attribute]
            gen_kwargs=dict(
                num_examples=self.builder_config.test_examples,  # pyrefly: ignore[missing-attribute]
                data_path=extracted["test_data"],
                label_path=extracted["test_labels"],
            ),
        ),
    ]


def _extract_mnist_images(image_filepath, num_images):
  with tf.io.gfile.GFile(image_filepath, "rb") as f:
    f.read(16)  # header
    buf = f.read(_MNIST_IMAGE_SIZE * _MNIST_IMAGE_SIZE * num_images)
    data = np.frombuffer(
        buf,
        dtype=np.uint8,
    ).reshape(num_images, _MNIST_IMAGE_SIZE, _MNIST_IMAGE_SIZE, 1)
    return data


def _extract_mnist_labels(labels_filepath, num_labels):
  with tf.io.gfile.GFile(labels_filepath, "rb") as f:
    f.read(8)  # header
    buf = f.read(num_labels)
    labels = np.frombuffer(buf, dtype=np.uint8).astype(np.int64)
    return labels