text: from .vision import VisionDataset
  import warnings
  from PIL import Image
  import os
  import os.path
  import numpy as np
  import torch
  import codecs
  import string
  from typing import Any, Callable, Dict, IO, List, Optional, Tuple, Union
  from .utils import download_url, download_and_extract_archive, extract_archive, \
      verify_str_arg
  
  
  class MNIST(VisionDataset):
      """`MNIST <http://yann.lecun.com/exdb/mnist/>`_ Dataset.
  
      Args:
          root (string): Root directory of dataset where ``MNIST/processed/training.pt``
              and  ``MNIST/processed/test.pt`` exist.
          train (bool, optional): If True, creates dataset from ``training.pt``,
              otherwise from ``test.pt``.
          download (bool, optional): If true, downloads the dataset from the internet and
              puts it in root directory. If dataset is already downloaded, it is not
              downloaded again.
          transform (callable, optional): A function/transform that  takes in an PIL image
              and returns a transformed version. E.g, ``transforms.RandomCrop``
          target_transform (callable, optional): A function/transform that takes in the
              target and transforms it.
      """
  
      resources = [
          ("http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz", "f68b3c2dcbeaaa9fbdd348bbdeb94873"),
          ("http://yann.lecun.com/exdb/mnist/train-labels-idx1-ubyte.gz", "d53e105ee54ea40749a09fcbcd1e9432"),
          ("http://yann.lecun.com/exdb/mnist/t10k-images-idx3-ubyte.gz", "9fb629c4189551a2d022fa330f9573f3"),
          ("http://yann.lecun.com/exdb/mnist/t10k-labels-idx1-ubyte.gz", "ec29112dd5afa0611ce80d1b7f02629c")
      ]
  
      training_file = 'training.pt'
      test_file = 'test.pt'
      classes = ['0 - zero', '1 - one', '2 - two', '3 - three', '4 - four',
                 '5 - five', '6 - six', '7 - seven', '8 - eight', '9 - nine']
  
      @property
      def train_labels(self):
          warnings.warn("train_labels has been renamed targets")
          return self.targets
  
      @property
      def test_labels(self):
          warnings.warn("test_labels has been renamed targets")
          return self.targets
  
      @property
      def train_data(self):
          warnings.warn("train_data has been renamed data")
          return self.data
  
      @property
      def test_data(self):
          warnings.warn("test_data has been renamed data")
          return self.data
  
      def __init__(
              self,
              root: str,
              train: bool = True,
              transform: Optional[Callable] = None,
              target_transform: Optional[Callable] = None,
              download: bool = False,
      ) -> None:
          super(MNIST, self).__init__(root, transform=transform,
                                      target_transform=target_transform)
          self.train = train  # training set or test set
  
          if download:
              self.download()
  
          if not self._check_exists():
              raise RuntimeError('Dataset not found.' +
                                 ' You can use download=True to download it')
  
          if self.train:
              data_file = self.training_file
          else:
              data_file = self.test_file
          self.data, self.targets = torch.load(os.path.join(self.processed_folder, data_file))
  
      def __getitem__(self, index: int) -> Tuple[Any, Any]:
          """
          Args:
              index (int): Index
  
          Returns:
              tuple: (image, target) where target is index of the target class.
          """
          img, target = self.data[index], int(self.targets[index])
  
          # doing this so that it is consistent with all other datasets
          # to return a PIL Image
          img = Image.fromarray(img.numpy(), mode='L')
  
          if self.transform is not None:
              img = self.transform(img)
  
          if self.target_transform is not None:
              target = self.target_transform(target)
  
          return img, target
  
      def __len__(self) -> int:
          return len(self.data)
  
      @property
      def raw_folder(self) -> str:
          return os.path.join(self.root, self.__class__.__name__, 'raw')
  
      @property
      def processed_folder(self) -> str:
          return os.path.join(self.root, self.__class__.__name__, 'processed')
  
      @property
      def class_to_idx(self) -> Dict[str, int]:
          return {_class: i for i, _class in enumerate(self.classes)}
  
      def _check_exists(self) -> bool:
          return (os.path.exists(os.path.join(self.processed_folder,
                                              self.training_file)) and
                  os.path.exists(os.path.join(self.processed_folder,
                                              self.test_file)))
  
      def download(self) -> None:
          """Download the MNIST data if it doesn't exist in processed_folder already."""
  
          if self._check_exists():
              return
  
          os.makedirs(self.raw_folder, exist_ok=True)
          os.makedirs(self.processed_folder, exist_ok=True)
  
          # download files
          for url, md5 in self.resources:
              filename = url.rpartition('/')[2]
              download_and_extract_archive(url, download_root=self.raw_folder, filename=filename, md5=md5)
  
          # process and save as torch files
          print('Processing...')
  
          training_set = (
              read_image_file(os.path.join(self.raw_folder, 'train-images-idx3-ubyte')),
              read_label_file(os.path.join(self.raw_folder, 'train-labels-idx1-ubyte'))
          )
          test_set = (
              read_image_file(os.path.join(self.raw_folder, 't10k-images-idx3-ubyte')),
              read_label_file(os.path.join(self.raw_folder, 't10k-labels-idx1-ubyte'))
          )
          with open(os.path.join(self.processed_folder, self.training_file), 'wb') as f:
              torch.save(training_set, f)
          with open(os.path.join(self.processed_folder, self.test_file), 'wb') as f:
              torch.save(test_set, f)
  
          print('Done!')
  
      def extra_repr(self) -> str:
          return "Split: {}".format("Train" if self.train is True else "Test")
  
  
  class FashionMNIST(MNIST):
      """`Fashion-MNIST <https://github.com/zalandoresearch/fashion-mnist>`_ Dataset.
  
      Args:
          root (string): Root directory of dataset where ``Fashion-MNIST/processed/training.pt``
              and  ``Fashion-MNIST/processed/test.pt`` exist.
          train (bool, optional): If True, creates dataset from ``training.pt``,
              otherwise from ``test.pt``.
          download (bool, optional): If true, downloads the dataset from the internet and
              puts it in root directory. If dataset is already downloaded, it is not
              downloaded again.
          transform (callable, optional): A function/transform that  takes in an PIL image
              and returns a transformed version. E.g, ``transforms.RandomCrop``
          target_transform (callable, optional): A function/transform that takes in the
              target and transforms it.
      """
      resources = [
          ("http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/train-images-idx3-ubyte.gz",
           "8d4fb7e6c68d591d4c3dfef9ec88bf0d"),
          ("http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/train-labels-idx1-ubyte.gz",
           "25c81989df183df01b3e8a0aad5dffbe"),
          ("http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/t10k-images-idx3-ubyte.gz",
           "bef4ecab320f06d8554ea6380940ec79"),
          ("http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/t10k-labels-idx1-ubyte.gz",
           "bb300cfdad3c16e7a12a480ee83cd310")
      ]
      classes = ['T-shirt/top', 'Trouser', 'Pullover', 'Dress', 'Coat', 'Sandal',
                 'Shirt', 'Sneaker', 'Bag', 'Ankle boot']
  
  
  class KMNIST(MNIST):
      """`Kuzushiji-MNIST <https://github.com/rois-codh/kmnist>`_ Dataset.
  
      Args:
          root (string): Root directory of dataset where ``KMNIST/processed/training.pt``
              and  ``KMNIST/processed/test.pt`` exist.
          train (bool, optional): If True, creates dataset from ``training.pt``,
              otherwise from ``test.pt``.
          download (bool, optional): If true, downloads the dataset from the internet and
              puts it in root directory. If dataset is already downloaded, it is not
              downloaded again.
          transform (callable, optional): A function/transform that  takes in an PIL image
              and returns a transformed version. E.g, ``transforms.RandomCrop``
          target_transform (callable, optional): A function/transform that takes in the
              target and transforms it.
      """
      resources = [
          ("http://codh.rois.ac.jp/kmnist/dataset/kmnist/train-images-idx3-ubyte.gz", "bdb82020997e1d708af4cf47b453dcf7"),
          ("http://codh.rois.ac.jp/kmnist/dataset/kmnist/train-labels-idx1-ubyte.gz", "e144d726b3acfaa3e44228e80efcd344"),
          ("http://codh.rois.ac.jp/kmnist/dataset/kmnist/t10k-images-idx3-ubyte.gz", "5c965bf0a639b31b8f53240b1b52f4d7"),
          ("http://codh.rois.ac.jp/kmnist/dataset/kmnist/t10k-labels-idx1-ubyte.gz", "7320c461ea6c1c855c0b718fb2a4b134")
      ]
      classes = ['o', 'ki', 'su', 'tsu', 'na', 'ha', 'ma', 'ya', 're', 'wo']
  
  
  class EMNIST(MNIST):
      """`EMNIST <https://www.westernsydney.edu.au/bens/home/reproducible_research/emnist>`_ Dataset.
  
      Args:
          root (string): Root directory of dataset where ``EMNIST/processed/training.pt``
              and  ``EMNIST/processed/test.pt`` exist.
          split (string): The dataset has 6 different splits: ``byclass``, ``bymerge``,
              ``balanced``, ``letters``, ``digits`` and ``mnist``. This argument specifies
              which one to use.
          train (bool, optional): If True, creates dataset from ``training.pt``,
              otherwise from ``test.pt``.
          download (bool, optional): If true, downloads the dataset from the internet and
              puts it in root directory. If dataset is already downloaded, it is not
              downloaded again.
          transform (callable, optional): A function/transform that  takes in an PIL image
              and returns a transformed version. E.g, ``transforms.RandomCrop``
          target_transform (callable, optional): A function/transform that takes in the
              target and transforms it.
      """
      # Updated URL from https://www.nist.gov/node/1298471/emnist-dataset since the
      # _official_ download link
      # https://cloudstor.aarnet.edu.au/plus/s/ZNmuFiuQTqZlu9W/download
      # is (currently) unavailable
      url = 'http://www.itl.nist.gov/iaui/vip/cs_links/EMNIST/gzip.zip'
      md5 = "58c8d27c78d21e728a6bc7b3cc06412e"
      splits = ('byclass', 'bymerge', 'balanced', 'letters', 'digits', 'mnist')
      # Merged Classes assumes Same structure for both uppercase and lowercase version
      _merged_classes = set(['C', 'I', 'J', 'K', 'L', 'M', 'O', 'P', 'S', 'U', 'V', 'W', 'X', 'Y', 'Z'])
      _all_classes = set(list(string.digits + string.ascii_letters))
      classes_split_dict = {
          'byclass': list(_all_classes),
          'bymerge': sorted(list(_all_classes - _merged_classes)),
          'balanced': sorted(list(_all_classes - _merged_classes)),
          'letters': list(string.ascii_lowercase),
          'digits': list(string.digits),
          'mnist': list(string.digits),
      }
  
      def __init__(self, root: str, split: str, **kwargs: Any) -> None:
          self.split = verify_str_arg(split, "split", self.splits)
          self.training_file = self._training_file(split)
          self.test_file = self._test_file(split)
          super(EMNIST, self).__init__(root, **kwargs)
          self.classes = self.classes_split_dict[self.split]
  
      @staticmethod
      def _training_file(split) -> str:
          return 'training_{}.pt'.format(split)
  
      @staticmethod
      def _test_file(split) -> str:
          return 'test_{}.pt'.format(split)
  
      def download(self) -> None:
          """Download the EMNIST data if it doesn't exist in processed_folder already."""
          import shutil
  
          if self._check_exists():
              return
  
          os.makedirs(self.raw_folder, exist_ok=True)
          os.makedirs(self.processed_folder, exist_ok=True)
  
          # download files
          print('Downloading and extracting zip archive')
          download_and_extract_archive(self.url, download_root=self.raw_folder, filename="emnist.zip",
                                       remove_finished=True, md5=self.md5)
          gzip_folder = os.path.join(self.raw_folder, 'gzip')
          for gzip_file in os.listdir(gzip_folder):
              if gzip_file.endswith('.gz'):
                  extract_archive(os.path.join(gzip_folder, gzip_file), gzip_folder)
  
          # process and save as torch files
          for split in self.splits:
              print('Processing ' + split)
              training_set = (
                  read_image_file(os.path.join(gzip_folder, 'emnist-{}-train-images-idx3-ubyte'.format(split))),
                  read_label_file(os.path.join(gzip_folder, 'emnist-{}-train-labels-idx1-ubyte'.format(split)))
              )
              test_set = (
                  read_image_file(os.path.join(gzip_folder, 'emnist-{}-test-images-idx3-ubyte'.format(split))),
                  read_label_file(os.path.join(gzip_folder, 'emnist-{}-test-labels-idx1-ubyte'.format(split)))
              )
              with open(os.path.join(self.processed_folder, self._training_file(split)), 'wb') as f:
                  torch.save(training_set, f)
              with open(os.path.join(self.processed_folder, self._test_file(split)), 'wb') as f:
                  torch.save(test_set, f)
          shutil.rmtree(gzip_folder)
  
          print('Done!')
  
  
  class QMNIST(MNIST):
      """`QMNIST <https://github.com/facebookresearch/qmnist>`_ Dataset.
  
      Args:
          root (string): Root directory of dataset whose ``processed''
              subdir contains torch binary files with the datasets.
          what (string,optional): Can be 'train', 'test', 'test10k',
              'test50k', or 'nist' for respectively the mnist compatible
              training set, the 60k qmnist testing set, the 10k qmnist
              examples that match the mnist testing set, the 50k
              remaining qmnist testing examples, or all the nist
              digits. The default is to select 'train' or 'test'
              according to the compatibility argument 'train'.
          compat (bool,optional): A boolean that says whether the target
              for each example is class number (for compatibility with
              the MNIST dataloader) or a torch vector containing the
              full qmnist information. Default=True.
          download (bool, optional): If true, downloads the dataset from
              the internet and puts it in root directory. If dataset is
              already downloaded, it is not downloaded again.
          transform (callable, optional): A function/transform that
              takes in an PIL image and returns a transformed
              version. E.g, ``transforms.RandomCrop``
          target_transform (callable, optional): A function/transform
              that takes in the target and transforms it.
          train (bool,optional,compatibility): When argument 'what' is
              not specified, this boolean decides whether to load the
              training set ot the testing set.  Default: True.
  
      """
  
      subsets = {
          'train': 'train',
          'test': 'test',
          'test10k': 'test',
          'test50k': 'test',
          'nist': 'nist'
      }
      resources: Dict[str, List[Tuple[str, str]]] = {  # type: ignore[assignment]
          'train': [('https://raw.githubusercontent.com/facebookresearch/qmnist/master/qmnist-train-images-idx3-ubyte.gz',
                     'ed72d4157d28c017586c42bc6afe6370'),
                    ('https://raw.githubusercontent.com/facebookresearch/qmnist/master/qmnist-train-labels-idx2-int.gz',
                     '0058f8dd561b90ffdd0f734c6a30e5e4')],
          'test': [('https://raw.githubusercontent.com/facebookresearch/qmnist/master/qmnist-test-images-idx3-ubyte.gz',
                    '1394631089c404de565df7b7aeaf9412'),
                   ('https://raw.githubusercontent.com/facebookresearch/qmnist/master/qmnist-test-labels-idx2-int.gz',
                    '5b5b05890a5e13444e108efe57b788aa')],
          'nist': [('https://raw.githubusercontent.com/facebookresearch/qmnist/master/xnist-images-idx3-ubyte.xz',
                    '7f124b3b8ab81486c9d8c2749c17f834'),
                   ('https://raw.githubusercontent.com/facebookresearch/qmnist/master/xnist-labels-idx2-int.xz',
                    '5ed0e788978e45d4a8bd4b7caec3d79d')]
      }
      classes = ['0 - zero', '1 - one', '2 - two', '3 - three', '4 - four',
                 '5 - five', '6 - six', '7 - seven', '8 - eight', '9 - nine']
  
      def __init__(
              self, root: str, what: Optional[str] = None, compat: bool = True,
              train: bool = True, **kwargs: Any
      ) -> None:
          if what is None:
              what = 'train' if train else 'test'
          self.what = verify_str_arg(what, "what", tuple(self.subsets.keys()))
          self.compat = compat
          self.data_file = what + '.pt'
          self.training_file = self.data_file
          self.test_file = self.data_file
          super(QMNIST, self).__init__(root, train, **kwargs)
  
      def download(self) -> None:
          """Download the QMNIST data if it doesn't exist in processed_folder already.
             Note that we only download what has been asked for (argument 'what').
          """
          if self._check_exists():
              return
          os.makedirs(self.raw_folder, exist_ok=True)
          os.makedirs(self.processed_folder, exist_ok=True)
          split = self.resources[self.subsets[self.what]]
          files = []
  
          # download data files if not already there
          for url, md5 in split:
              filename = url.rpartition('/')[2]
              file_path = os.path.join(self.raw_folder, filename)
              if not os.path.isfile(file_path):
                  download_url(url, root=self.raw_folder, filename=filename, md5=md5)
              files.append(file_path)
  
          # process and save as torch files
          print('Processing...')
          data = read_sn3_pascalvincent_tensor(files[0])
          assert(data.dtype == torch.uint8)
          assert(data.ndimension() == 3)
          targets = read_sn3_pascalvincent_tensor(files[1]).long()
          assert(targets.ndimension() == 2)
          if self.what == 'test10k':
              data = data[0:10000, :, :].clone()
              targets = targets[0:10000, :].clone()
          if self.what == 'test50k':
              data = data[10000:, :, :].clone()
              targets = targets[10000:, :].clone()
          with open(os.path.join(self.processed_folder, self.data_file), 'wb') as f:
              torch.save((data, targets), f)
  
      def __getitem__(self, index: int) -> Tuple[Any, Any]:
          # redefined to handle the compat flag
          img, target = self.data[index], self.targets[index]
          img = Image.fromarray(img.numpy(), mode='L')
          if self.transform is not None:
              img = self.transform(img)
          if self.compat:
              target = int(target[0])
          if self.target_transform is not None:
              target = self.target_transform(target)
          return img, target
  
      def extra_repr(self) -> str:
          return "Split: {}".format(self.what)
  
  
  def get_int(b: bytes) -> int:
      return int(codecs.encode(b, 'hex'), 16)
  
  
  def open_maybe_compressed_file(path: Union[str, IO]) -> IO:
      """Return a file object that possibly decompresses 'path' on the fly.
         Decompression occurs when argument `path` is a string and ends with '.gz' or '.xz'.
      """
      if not isinstance(path, torch._six.string_classes):
          return path
      if path.endswith('.gz'):
          import gzip
          return gzip.open(path, 'rb')
      if path.endswith('.xz'):
          import lzma
          return lzma.open(path, 'rb')
      return open(path, 'rb')
  
  
  SN3_PASCALVINCENT_TYPEMAP = {
      8: (torch.uint8, np.uint8, np.uint8),
      9: (torch.int8, np.int8, np.int8),
      11: (torch.int16, np.dtype('>i2'), 'i2'),
      12: (torch.int32, np.dtype('>i4'), 'i4'),
      13: (torch.float32, np.dtype('>f4'), 'f4'),
      14: (torch.float64, np.dtype('>f8'), 'f8')
  }
  
  
  def read_sn3_pascalvincent_tensor(path: Union[str, IO], strict: bool = True) -> torch.Tensor:
      """Read a SN3 file in "Pascal Vincent" format (Lush file 'libidx/idx-io.lsh').
         Argument may be a filename, compressed filename, or file object.
      """
      # read
      with open_maybe_compressed_file(path) as f:
          data = f.read()
      # parse
      magic = get_int(data[0:4])
      nd = magic % 256
      ty = magic // 256
      assert nd >= 1 and nd <= 3
      assert ty >= 8 and ty <= 14
      m = SN3_PASCALVINCENT_TYPEMAP[ty]
      s = [get_int(data[4 * (i + 1): 4 * (i + 2)]) for i in range(nd)]
      parsed = np.frombuffer(data, dtype=m[1], offset=(4 * (nd + 1)))
      assert parsed.shape[0] == np.prod(s) or not strict
      return torch.from_numpy(parsed.astype(m[2], copy=False)).view(*s)
  
  
  def read_label_file(path: str) -> torch.Tensor:
      with open(path, 'rb') as f:
          x = read_sn3_pascalvincent_tensor(f, strict=False)
      assert(x.dtype == torch.uint8)
      assert(x.ndimension() == 1)
      return x.long()
  
  
  def read_image_file(path: str) -> torch.Tensor:
      with open(path, 'rb') as f:
          x = read_sn3_pascalvincent_tensor(f, strict=False)
      assert(x.dtype == torch.uint8)
      assert(x.ndimension() == 3)
      return x
  