text: r"""
  A *Reader* is a PyTorch :class:`~torch.utils.data.Dataset` which simply reads
  data from disk and returns it almost as is. Readers defined here are used by
  datasets in :mod:`virtex.data.datasets`.
  """
  from collections import defaultdict
  import glob
  import json
  import os
  import pickle
  import random
  from typing import Dict, List, Tuple
  
  import cv2
  import lmdb
  from loguru import logger
  from torch.utils.data import Dataset
  
  
  # Some simplified type renaming for better readability
  ImageID = int
  Captions = List[str]
  
  
  class SimpleCocoCaptionsReader(Dataset):
      r"""
      A reader interface to read COCO Captions dataset and directly from official
      annotation files and return it unprocessed. We only use this for serializing
      the dataset to LMDB files, and use :class:`~virtex.data.readers.LmdbReader`
      in rest of the datasets.
  
      Parameters
      ----------
      root: str, optional (default = "datasets/coco")
          Path to the COCO dataset root directory.
      split: str, optional (default = "train")
          Which split (from COCO 2017 version) to read. One of ``{"train", "val"}``.
      """
      def __init__(self, root: str = "datasets/coco", split: str = "train"):
  
          image_dir = os.path.join(root, f"{split}2017")
  
          # Make a tuple of image id and its filename, get image_id from its
          # filename (assuming directory has images with names in COCO2017 format).
          #image_filenames = glob.glob(os.path.join(image_dir, "*.jpg"))
          #self.id_filename: List[Tuple[ImageID, str]] = [
          #    (int(os.path.basename(name)[:-4]), name) for name in image_filenames
          #]
  
          # Make a mapping between image_id and its captions.
          _captions = json.load(
              open(os.path.join(root, "annotations", f"captions_{split}2017.json"))
          )
  
          self.id_filename: List[Tuple[ImageID, str]] = [
              (img['id'], os.path.join(image_dir, img['file_name'])) for img in _captions['images']
          ]
  
          self._id_to_captions: Dict[ImageID, Captions] = defaultdict(list)
  
          for ann in _captions["annotations"]:
              self._id_to_captions[ann["image_id"]].append(ann["caption"])
  
      def __len__(self):
          return len(self.id_filename)
  
      def __getitem__(self, idx: int):
          image_id, filename = self.id_filename[idx]
  
          try:
              # shape: (height, width, channels), dtype: uint8
              image = cv2.imread(filename)
              image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
              captions = self._id_to_captions[image_id]
          except:
              return {"image_id": image_id, 'filename': filename}
  
          return {"image_id": image_id, "image": image, "captions": captions}
  
  
  class LmdbReader(Dataset):
      r"""
      A reader interface to read datapoints from a serialized LMDB file containing
      ``(image_id, image, caption)`` tuples. Optionally, one may specify a
      partial percentage of datapoints to use.
  
      .. note::
  
          When training in distributed setting, make sure each worker has SAME
          random seed because there is some randomness in selecting keys for
          training with partial dataset. If you wish to use a different seed for
          each worker, select keys manually outside of this class and use
          :meth:`set_keys`.
  
      .. note::
  
          Similar to :class:`~torch.utils.data.distributed.DistributedSampler`,
          this reader can shuffle the dataset deterministically at the start of
          epoch. Use :meth:`set_shuffle_seed` manually from outside to change the
          seed at every epoch.
  
      Parameters
      ----------
      lmdb_path: str
          Path to LMDB file with datapoints.
      shuffle: bool, optional (default = True)
          Whether to shuffle or not. If this is on, there will be one deterministic
          shuffle based on epoch before sharding the dataset (to workers).
      percentage: float, optional (default = 100.0)
          Percentage of datapoints to use. If less than 100.0, keys will be
          shuffled and first K% will be retained and use throughout training.
          Make sure to set this only for training, not validation.
      """
  
      def __init__(self, lmdb_path: str, shuffle: bool = True, percentage: float = 100):
          self.lmdb_path = lmdb_path
          self.shuffle = shuffle
  
          assert percentage > 0, "Cannot load dataset with 0 percent original size."
          self.percentage = percentage
  
          # fmt: off
          # Create an LMDB transaction right here. It will be aborted when this
          # class goes out of scope.
          env = lmdb.open(
              self.lmdb_path, subdir=False, readonly=True, lock=False,
              readahead=False, map_size=1099511627776 * 2,
          )
          self.db_txn = env.begin()
  
          # Form a list of LMDB keys numbered from 0 (as binary strings).
          self._keys = [
              f"{i}".encode("ascii") for i in range(env.stat()["entries"])
          ]
          # fmt: on
  
          # If data percentage < 100%, randomly retain K% keys. This will be
          # deterministic based on random seed.
          if percentage < 100.0:
              retain_k: int = int(len(self._keys) * percentage / 100.0)
              random.shuffle(self._keys)
              self._keys = self._keys[:retain_k]
              logger.info(f"Retained {retain_k} datapoints for training!")
  
          # A seed to deterministically shuffle at the start of epoch. This is
          # set externally through `set_shuffle_seed`.
          self.shuffle_seed = 0
  
      def set_shuffle_seed(self, seed: int):
          r"""Set random seed for shuffling data."""
          self.shuffle_seed = seed
  
      def get_keys(self) -> List[bytes]:
          r"""Return list of keys, useful while saving checkpoint."""
          return self._keys
  
      def set_keys(self, keys: List[bytes]):
          r"""Set list of keys, useful while loading from checkpoint."""
          self._keys = keys
  
      def __getstate__(self):
          r"""
          This magic method allows an object of this class to be pickable, useful
          for dataloading with multiple CPU workers. :attr:`db_txn` is not
          pickable, so we remove it from state, and re-instantiate it in
          :meth:`__setstate__`.
          """
          state = self.__dict__
          state["db_txn"] = None
          return state
  
      def __setstate__(self, state):
          self.__dict__ = state
  
          env = lmdb.open(
              self.lmdb_path, subdir=False, readonly=True, lock=False,
              readahead=False, map_size=1099511627776 * 2,
          )
          self.db_txn = env.begin()
  
      def __len__(self):
          return len(self._keys)
  
      def __getitem__(self, idx: int):
          datapoint_pickled = self.db_txn.get(self._keys[idx])
          image_id, image, captions = pickle.loads(datapoint_pickled)
  
          return image_id, image, captions
  