text: from math import ceil, floor
  from typing import Dict, List, Optional
  
  import numpy as np
  import pytorch_lightning as pl
  import torch
  from torch.utils.data import DataLoader, Dataset
  
  from scvi import REGISTRY_KEYS, settings
  from scvi.data.anndata import AnnDataManager
  from scvi.dataloaders._ann_dataloader import AnnDataLoader, BatchSampler
  from scvi.dataloaders._semi_dataloader import SemiSupervisedDataLoader
  from scvi.model._utils import parse_use_gpu_arg
  
  
  def validate_data_split(
      n_samples: int, train_size: float, validation_size: Optional[float] = None
  ):
      """
      Check data splitting parameters and return n_train and n_val.
  
      Parameters
      ----------
      n_samples
          Number of samples to split
      train_size
          Size of train set. Need to be: 0 < train_size <= 1.
      validation_size
          Size of validation set. Need to be 0 <= validation_size < 1
      """
      if train_size > 1.0 or train_size <= 0.0:
          raise ValueError("Invalid train_size. Must be: 0 < train_size <= 1")
  
      n_train = ceil(train_size * n_samples)
  
      if validation_size is None:
          n_val = n_samples - n_train
      elif validation_size >= 1.0 or validation_size < 0.0:
          raise ValueError("Invalid validation_size. Must be 0 <= validation_size < 1")
      elif (train_size + validation_size) > 1:
          raise ValueError("train_size + validation_size must be between 0 and 1")
      else:
          n_val = floor(n_samples * validation_size)
  
      if n_train == 0:
          raise ValueError(
              "With n_samples={}, train_size={} and validation_size={}, the "
              "resulting train set will be empty. Adjust any of the "
              "aforementioned parameters.".format(n_samples, train_size, validation_size)
          )
  
      return n_train, n_val
  
  
  class DataSplitter(pl.LightningDataModule):
      """
      Creates data loaders ``train_set``, ``validation_set``, ``test_set``.
  
      If ``train_size + validation_set < 1`` then ``test_set`` is non-empty.
  
      Parameters
      ----------
      adata_manager
          :class:`~scvi.data.anndata.AnnDataManager` object that has been created via ``setup_anndata``.
      train_size
          float, or None (default is 0.9)
      validation_size
          float, or None (default is None)
      use_gpu
          Use default GPU if available (if None or True), or index of GPU to use (if int),
          or name of GPU (if str, e.g., `'cuda:0'`), or use CPU (if False).
      **kwargs
          Keyword args for data loader. If adata has labeled data, data loader
          class is :class:`~scvi.dataloaders.SemiSupervisedDataLoader`,
          else data loader class is :class:`~scvi.dataloaders.AnnDataLoader`.
  
      Examples
      --------
      >>> adata = scvi.data.synthetic_iid()
      >>> scvi.model.SCVI.setup_anndata(adata)
      >>> adata_manager = scvi.model.SCVI(adata).adata_manager
      >>> splitter = DataSplitter(adata)
      >>> splitter.setup()
      >>> train_dl = splitter.train_dataloader()
      """
  
      def __init__(
          self,
          adata_manager: AnnDataManager,
          train_size: float = 0.9,
          validation_size: Optional[float] = None,
          use_gpu: bool = False,
          **kwargs,
      ):
          super().__init__()
          self.adata_manager = adata_manager
          self.train_size = float(train_size)
          self.validation_size = validation_size
          self.data_loader_kwargs = kwargs
          self.use_gpu = use_gpu
  
          self.n_train, self.n_val = validate_data_split(
              self.adata_manager.adata.n_obs, self.train_size, self.validation_size
          )
  
      def setup(self, stage: Optional[str] = None):
          """Split indices in train/test/val sets."""
          n_train = self.n_train
          n_val = self.n_val
          random_state = np.random.RandomState(seed=settings.seed)
          permutation = random_state.permutation(self.adata_manager.adata.n_obs)
          self.val_idx = permutation[:n_val]
          self.train_idx = permutation[n_val : (n_val + n_train)]
          self.test_idx = permutation[(n_val + n_train) :]
  
          gpus, self.device = parse_use_gpu_arg(self.use_gpu, return_device=True)
          self.pin_memory = (
              True if (settings.dl_pin_memory_gpu_training and gpus != 0) else False
          )
  
      def train_dataloader(self):
          return AnnDataLoader(
              self.adata_manager,
              indices=self.train_idx,
              shuffle=True,
              drop_last=3,
              pin_memory=self.pin_memory,
              **self.data_loader_kwargs,
          )
  
      def val_dataloader(self):
          if len(self.val_idx) > 0:
              return AnnDataLoader(
                  self.adata_manager,
                  indices=self.val_idx,
                  shuffle=False,
                  drop_last=3,
                  pin_memory=self.pin_memory,
                  **self.data_loader_kwargs,
              )
          else:
              pass
  
      def test_dataloader(self):
          if len(self.test_idx) > 0:
              return AnnDataLoader(
                  self.adata_manager,
                  indices=self.test_idx,
                  shuffle=False,
                  drop_last=3,
                  pin_memory=self.pin_memory,
                  **self.data_loader_kwargs,
              )
          else:
              pass
  
  
  class SemiSupervisedDataSplitter(pl.LightningDataModule):
      """
      Creates data loaders ``train_set``, ``validation_set``, ``test_set``.
  
      If ``train_size + validation_set < 1`` then ``test_set`` is non-empty.
      The ratio between labeled and unlabeled data in adata will be preserved
      in the train/test/val sets.
  
      Parameters
      ----------
      adata_manager
          :class:`~scvi.data.anndata.AnnDataManager` object that has been created via ``setup_anndata``.
      unlabeled_category
          Category to treat as unlabeled
      train_size
          float, or None (default is 0.9)
      validation_size
          float, or None (default is None)
      n_samples_per_label
          Number of subsamples for each label class to sample per epoch
      use_gpu
          Use default GPU if available (if None or True), or index of GPU to use (if int),
          or name of GPU (if str, e.g., `'cuda:0'`), or use CPU (if False).
      **kwargs
          Keyword args for data loader. If adata has labeled data, data loader
          class is :class:`~scvi.dataloaders.SemiSupervisedDataLoader`,
          else data loader class is :class:`~scvi.dataloaders.AnnDataLoader`.
  
      Examples
      --------
      >>> adata = scvi.data.synthetic_iid()
      >>> scvi.model.SCVI.setup_anndata(adata, labels_key="labels")
      >>> adata_manager = scvi.model.SCVI(adata).adata_manager
      >>> unknown_label = 'label_0'
      >>> splitter = SemiSupervisedDataSplitter(adata, unknown_label)
      >>> splitter.setup()
      >>> train_dl = splitter.train_dataloader()
      """
  
      def __init__(
          self,
          adata_manager: AnnDataManager,
          unlabeled_category,
          train_size: float = 0.9,
          validation_size: Optional[float] = None,
          n_samples_per_label: Optional[int] = None,
          use_gpu: bool = False,
          **kwargs,
      ):
          super().__init__()
          self.adata_manager = adata_manager
          self.unlabeled_category = unlabeled_category
          self.train_size = float(train_size)
          self.validation_size = validation_size
          self.data_loader_kwargs = kwargs
          self.n_samples_per_label = n_samples_per_label
  
          original_key = adata_manager.get_state_registry(
              REGISTRY_KEYS.LABELS_KEY
          ).original_key
          labels = np.asarray(adata_manager.adata.obs[original_key]).ravel()
          self._unlabeled_indices = np.argwhere(labels == unlabeled_category).ravel()
          self._labeled_indices = np.argwhere(labels != unlabeled_category).ravel()
  
          self.data_loader_kwargs = kwargs
          self.use_gpu = use_gpu
  
      def setup(self, stage: Optional[str] = None):
          """Split indices in train/test/val sets."""
          n_labeled_idx = len(self._labeled_indices)
          n_unlabeled_idx = len(self._unlabeled_indices)
  
          if n_labeled_idx != 0:
              n_labeled_train, n_labeled_val = validate_data_split(
                  n_labeled_idx, self.train_size, self.validation_size
              )
              rs = np.random.RandomState(seed=settings.seed)
              labeled_permutation = rs.choice(
                  self._labeled_indices, len(self._labeled_indices), replace=False
              )
              labeled_idx_val = labeled_permutation[:n_labeled_val]
              labeled_idx_train = labeled_permutation[
                  n_labeled_val : (n_labeled_val + n_labeled_train)
              ]
              labeled_idx_test = labeled_permutation[(n_labeled_val + n_labeled_train) :]
          else:
              labeled_idx_test = []
              labeled_idx_train = []
              labeled_idx_val = []
  
          if n_unlabeled_idx != 0:
              n_unlabeled_train, n_unlabeled_val = validate_data_split(
                  n_unlabeled_idx, self.train_size, self.validation_size
              )
              rs = np.random.RandomState(seed=settings.seed)
              unlabeled_permutation = rs.choice(
                  self._unlabeled_indices, len(self._unlabeled_indices)
              )
              unlabeled_idx_val = unlabeled_permutation[:n_unlabeled_val]
              unlabeled_idx_train = unlabeled_permutation[
                  n_unlabeled_val : (n_unlabeled_val + n_unlabeled_train)
              ]
              unlabeled_idx_test = unlabeled_permutation[
                  (n_unlabeled_val + n_unlabeled_train) :
              ]
          else:
              unlabeled_idx_train = []
              unlabeled_idx_val = []
              unlabeled_idx_test = []
  
          indices_train = np.concatenate((labeled_idx_train, unlabeled_idx_train))
          indices_val = np.concatenate((labeled_idx_val, unlabeled_idx_val))
          indices_test = np.concatenate((labeled_idx_test, unlabeled_idx_test))
  
          self.train_idx = indices_train.astype(int)
          self.val_idx = indices_val.astype(int)
          self.test_idx = indices_test.astype(int)
  
          gpus = parse_use_gpu_arg(self.use_gpu, return_device=False)
          self.pin_memory = (
              True if (settings.dl_pin_memory_gpu_training and gpus != 0) else False
          )
  
          if len(self._labeled_indices) != 0:
              self.data_loader_class = SemiSupervisedDataLoader
              dl_kwargs = {
                  "unlabeled_category": self.unlabeled_category,
                  "n_samples_per_label": self.n_samples_per_label,
              }
          else:
              self.data_loader_class = AnnDataLoader
              dl_kwargs = {}
  
          self.data_loader_kwargs.update(dl_kwargs)
  
      def train_dataloader(self):
          return self.data_loader_class(
              self.adata_manager,
              indices=self.train_idx,
              shuffle=True,
              drop_last=3,
              pin_memory=self.pin_memory,
              **self.data_loader_kwargs,
          )
  
      def val_dataloader(self):
          if len(self.val_idx) > 0:
              return self.data_loader_class(
                  self.adata_manager,
                  indices=self.val_idx,
                  shuffle=False,
                  drop_last=3,
                  pin_memory=self.pin_memory,
                  **self.data_loader_kwargs,
              )
          else:
              pass
  
      def test_dataloader(self):
          if len(self.test_idx) > 0:
              return self.data_loader_class(
                  self.adata_manager,
                  indices=self.test_idx,
                  shuffle=False,
                  drop_last=3,
                  pin_memory=self.pin_memory,
                  **self.data_loader_kwargs,
              )
          else:
              pass
  
  
  class DeviceBackedDataSplitter(DataSplitter):
      """
      Creates loaders for data that is already on device, e.g., GPU.
  
      If ``train_size + validation_set < 1`` then ``test_set`` is non-empty.
  
      Parameters
      ----------
      adata_manager
          :class:`~scvi.data.anndata.AnnDataManager` object that has been created via ``setup_anndata``.
      train_size
          float, or None (default is 0.9)
      validation_size
          float, or None (default is None)
      use_gpu
          Use default GPU if available (if None or True), or index of GPU to use (if int),
          or name of GPU (if str, e.g., `'cuda:0'`), or use CPU (if False).
      shuffle
          if ``True``, shuffles indices before sampling for training set
      shuffle_test_val
          Shuffle test and validation indices.
      batch_size
          batch size of each iteration. If `None`, do not minibatch
  
      Examples
      --------
      >>> adata = scvi.data.synthetic_iid()
      >>> scvi.model.SCVI.setup_anndata(adata)
      >>> adata_manager = scvi.model.SCVI(adata).adata_manager
      >>> splitter = DeviceBackedDataSplitter(adata)
      >>> splitter.setup()
      >>> train_dl = splitter.train_dataloader()
      """
  
      def __init__(
          self,
          adata_manager: AnnDataManager,
          train_size: float = 1.0,
          validation_size: Optional[float] = None,
          use_gpu: bool = False,
          shuffle: bool = False,
          shuffle_test_val: bool = False,
          batch_size: Optional[int] = None,
          **kwargs,
      ):
          super().__init__(
              adata_manager=adata_manager,
              train_size=train_size,
              validation_size=validation_size,
              use_gpu=use_gpu,
              **kwargs,
          )
          self.batch_size = batch_size
          self.shuffle = shuffle
          self.shuffle_test_val = shuffle_test_val
  
      def setup(self, stage: Optional[str] = None):
          super().setup()
  
          if self.shuffle is False:
              self.train_idx = np.sort(self.train_idx)
              self.val_idx = (
                  np.sort(self.val_idx) if len(self.val_idx) > 0 else self.val_idx
              )
              self.test_idx = (
                  np.sort(self.test_idx) if len(self.test_idx) > 0 else self.test_idx
              )
  
          self.train_tensor_dict = self._get_tensor_dict(
              self.train_idx, device=self.device
          )
          self.test_tensor_dict = self._get_tensor_dict(self.test_idx, device=self.device)
          self.val_tensor_dict = self._get_tensor_dict(self.val_idx, device=self.device)
  
      def _get_tensor_dict(self, indices, device):
          if len(indices) is not None and len(indices) > 0:
              dl = AnnDataLoader(
                  self.adata_manager,
                  indices=indices,
                  batch_size=len(indices),
                  shuffle=False,
                  pin_memory=self.pin_memory,
                  **self.data_loader_kwargs,
              )
              # will only have one minibatch
              for batch in dl:
                  tensor_dict = batch
  
              for k, v in tensor_dict.items():
                  tensor_dict[k] = v.to(device)
  
              return tensor_dict
          else:
              return None
  
      def _make_dataloader(self, tensor_dict: Dict[str, torch.Tensor], shuffle):
          if tensor_dict is None:
              return None
          dataset = _DeviceBackedDataset(tensor_dict)
          indices = np.arange(len(dataset))
          bs = self.batch_size if self.batch_size is not None else len(indices)
          sampler = BatchSampler(shuffle=shuffle, indices=indices, batch_size=bs)
          return DataLoader(dataset, sampler=sampler, batch_size=None)
  
      def train_dataloader(self):
          return self._make_dataloader(self.train_tensor_dict, self.shuffle)
  
      def test_dataloader(self):
          return self._make_dataloader(self.test_tensor_dict, self.shuffle_test_val)
  
      def val_dataloader(self):
          return self._make_dataloader(self.val_tensor_dict, self.shuffle_test_val)
  
  
  class _DeviceBackedDataset(Dataset):
      def __init__(self, tensor_dict: Dict[str, torch.Tensor]):
          self.data = tensor_dict
  
      def __getitem__(self, idx: List[int]) -> Dict[str, torch.Tensor]:
          return_dict = {}
          for key, value in self.data.items():
              return_dict[key] = value[idx]
  
          return return_dict
  
      def __len__(self):
          for _, value in self.data.items():
              return len(value)
  