text: import logging
  import warnings
  from typing import Optional, Sequence, Union
  
  import numpy as np
  import pandas as pd
  import torch
  from anndata import AnnData
  from pandas.api.types import CategoricalDtype
  
  from scvi import _CONSTANTS
  from scvi._compat import Literal
  from scvi.data._anndata import _make_obs_column_categorical
  from scvi.dataloaders import (
      AnnDataLoader,
      SemiSupervisedDataLoader,
      SemiSupervisedDataSplitter,
  )
  from scvi.module import SCANVAE
  from scvi.train import SemiSupervisedTrainingPlan, TrainRunner
  from scvi.train._callbacks import SubSampleLabels
  
  from ._scvi import SCVI
  from .base import ArchesMixin, BaseModelClass, RNASeqMixin, VAEMixin
  
  logger = logging.getLogger(__name__)
  
  
  class SCANVI(RNASeqMixin, VAEMixin, ArchesMixin, BaseModelClass):
      """
      Single-cell annotation using variational inference [Xu21]_.
  
      Inspired from M1 + M2 model, as described in (https://arxiv.org/pdf/1406.5298.pdf).
  
      Parameters
      ----------
      adata
          AnnData object that has been registered via :func:`~scvi.data.setup_anndata`.
      unlabeled_category
          Value used for unlabeled cells in `labels_key` used to setup AnnData with scvi.
      n_hidden
          Number of nodes per hidden layer.
      n_latent
          Dimensionality of the latent space.
      n_layers
          Number of hidden layers used for encoder and decoder NNs.
      dropout_rate
          Dropout rate for neural networks.
      dispersion
          One of the following:
  
          * ``'gene'`` - dispersion parameter of NB is constant per gene across cells
          * ``'gene-batch'`` - dispersion can differ between different batches
          * ``'gene-label'`` - dispersion can differ between different labels
          * ``'gene-cell'`` - dispersion can differ for every gene in every cell
      gene_likelihood
          One of:
  
          * ``'nb'`` - Negative binomial distribution
          * ``'zinb'`` - Zero-inflated negative binomial distribution
          * ``'poisson'`` - Poisson distribution
      **model_kwargs
          Keyword args for :class:`~scvi.module.SCANVAE`
  
      Examples
      --------
      >>> adata = anndata.read_h5ad(path_to_anndata)
      >>> scvi.data.setup_anndata(adata, batch_key="batch", labels_key="labels")
      >>> vae = scvi.model.SCANVI(adata, "Unknown")
      >>> vae.train()
      >>> adata.obsm["X_scVI"] = vae.get_latent_representation()
      >>> adata.obs["pred_label"] = vae.predict()
  
      Notes
      -----
      See further usage examples in the following tutorials:
  
      1. :doc:`/user_guide/notebooks/harmonization`
      2. :doc:`/user_guide/notebooks/scarches_scvi_tools`
      3. :doc:`/user_guide/notebooks/seed_labeling`
      """
  
      def __init__(
          self,
          adata: AnnData,
          unlabeled_category: Union[str, int, float],
          n_hidden: int = 128,
          n_latent: int = 10,
          n_layers: int = 1,
          dropout_rate: float = 0.1,
          dispersion: Literal["gene", "gene-batch", "gene-label", "gene-cell"] = "gene",
          gene_likelihood: Literal["zinb", "nb", "poisson"] = "zinb",
          **model_kwargs,
      ):
          super(SCANVI, self).__init__(adata)
          scanvae_model_kwargs = dict(model_kwargs)
  
          self.unlabeled_category_ = unlabeled_category
          has_unlabeled = self._set_indices_and_labels()
  
          if len(self._labeled_indices) != 0:
              self._dl_cls = SemiSupervisedDataLoader
          else:
              self._dl_cls = AnnDataLoader
  
          # ignores unlabeled catgegory
          n_labels = (
              self.summary_stats["n_labels"] - 1
              if has_unlabeled
              else self.summary_stats["n_labels"]
          )
          n_cats_per_cov = (
              self.scvi_setup_dict_["extra_categoricals"]["n_cats_per_key"]
              if "extra_categoricals" in self.scvi_setup_dict_
              else None
          )
          self.module = SCANVAE(
              n_input=self.summary_stats["n_vars"],
              n_batch=self.summary_stats["n_batch"],
              n_labels=n_labels,
              n_continuous_cov=self.summary_stats["n_continuous_covs"],
              n_cats_per_cov=n_cats_per_cov,
              n_hidden=n_hidden,
              n_latent=n_latent,
              n_layers=n_layers,
              dropout_rate=dropout_rate,
              dispersion=dispersion,
              gene_likelihood=gene_likelihood,
              **scanvae_model_kwargs,
          )
  
          self.unsupervised_history_ = None
          self.semisupervised_history_ = None
  
          self._model_summary_string = (
              "ScanVI Model with the following params: \nunlabeled_category: {}, n_hidden: {}, n_latent: {}"
              ", n_layers: {}, dropout_rate: {}, dispersion: {}, gene_likelihood: {}"
          ).format(
              unlabeled_category,
              n_hidden,
              n_latent,
              n_layers,
              dropout_rate,
              dispersion,
              gene_likelihood,
          )
          self.init_params_ = self._get_init_params(locals())
          self.was_pretrained = False
  
      @classmethod
      def from_scvi_model(
          cls,
          scvi_model: SCVI,
          unlabeled_category: str,
          adata: Optional[AnnData] = None,
          **scanvi_kwargs,
      ):
          """
          Initialize scanVI model with weights from pretrained scVI model.
  
          Parameters
          ----------
          scvi_model
              Pretrained scvi model
          unlabeled_category
              Value used for unlabeled cells in `labels_key` used to setup AnnData with scvi.
          adata
              AnnData object that has been registered via :func:`~scvi.data.setup_anndata`.
          scanvi_kwargs
              kwargs for scanVI model
          """
          if scvi_model.is_trained_ is False:
              warnings.warn("Passed in scvi model hasn't been trained yet.")
  
          scanvi_kwargs = dict(scanvi_kwargs)
          init_params = scvi_model.init_params_
          non_kwargs = init_params["non_kwargs"]
          kwargs = init_params["kwargs"]
          kwargs = {k: v for (i, j) in kwargs.items() for (k, v) in j.items()}
          for k, v in {**non_kwargs, **kwargs}.items():
              if k in scanvi_kwargs.keys():
                  warnings.warn(
                      "Ignoring param '{}' as it was already passed in to ".format(k)
                      + "pretrained scvi model with value {}.".format(v)
                  )
                  del scanvi_kwargs[k]
  
          if adata is None:
              adata = scvi_model.adata
  
          scanvi_model = cls(
              adata, unlabeled_category, **non_kwargs, **kwargs, **scanvi_kwargs
          )
          scvi_state_dict = scvi_model.module.state_dict()
          scanvi_model.module.load_state_dict(scvi_state_dict, strict=False)
          scanvi_model.was_pretrained = True
  
          return scanvi_model
  
      def _set_indices_and_labels(self):
          """
          Set indices and make unlabeled cat as the last cat.
  
          Returns
          -------
          True is categories reordered else False
          """
          # get indices for labeled and unlabeled cells
          key = self.scvi_setup_dict_["data_registry"][_CONSTANTS.LABELS_KEY]["attr_key"]
          mapping = self.scvi_setup_dict_["categorical_mappings"][key]["mapping"]
          original_key = self.scvi_setup_dict_["categorical_mappings"][key][
              "original_key"
          ]
          labels = np.asarray(self.adata.obs[original_key]).ravel()
  
          if self.unlabeled_category_ in labels:
              unlabeled_idx = np.where(mapping == self.unlabeled_category_)
              unlabeled_idx = unlabeled_idx[0][0]
              # move unlabeled category to be the last position
              mapping[unlabeled_idx], mapping[-1] = mapping[-1], mapping[unlabeled_idx]
              cat_dtype = CategoricalDtype(categories=mapping, ordered=True)
              # rerun setup for the batch column
              _make_obs_column_categorical(
                  self.adata,
                  original_key,
                  "_scvi_labels",
                  categorical_dtype=cat_dtype,
              )
              remapped = True
          else:
              remapped = False
  
          self.scvi_setup_dict_ = self.adata.uns["_scvi"]
          self._label_mapping = mapping
          # set unlabeled and labeled indices
          self._unlabeled_indices = np.argwhere(
              labels == self.unlabeled_category_
          ).ravel()
          self._labeled_indices = np.argwhere(labels != self.unlabeled_category_).ravel()
          self._code_to_label = {i: l for i, l in enumerate(self._label_mapping)}
          self.original_label_key = original_key
  
          return remapped
  
      def predict(
          self,
          adata: Optional[AnnData] = None,
          indices: Optional[Sequence[int]] = None,
          soft: bool = False,
          batch_size: Optional[int] = None,
      ) -> Union[np.ndarray, pd.DataFrame]:
          """
          Return cell label predictions.
  
          Parameters
          ----------
          adata
              AnnData object that has been registered via :func:`~scvi.data.setup_anndata`.
          indices
              Return probabilities for each class label.
          soft
              If True, returns per class probabilities
          batch_size
              Minibatch size for data loading into model. Defaults to `scvi.settings.batch_size`.
          """
          adata = self._validate_anndata(adata)
  
          if indices is None:
              indices = np.arange(adata.n_obs)
  
          scdl = self._make_data_loader(
              adata=adata,
              indices=indices,
              batch_size=batch_size,
          )
          y_pred = []
          for _, tensors in enumerate(scdl):
              x = tensors[_CONSTANTS.X_KEY]
              batch = tensors[_CONSTANTS.BATCH_KEY]
              pred = self.module.classify(x, batch)
              if not soft:
                  pred = pred.argmax(dim=1)
              y_pred.append(pred.detach().cpu())
  
          y_pred = torch.cat(y_pred).numpy()
          if not soft:
              predictions = []
              for p in y_pred:
                  predictions.append(self._code_to_label[p])
  
              return np.array(predictions)
          else:
              n_labels = len(pred[0])
              pred = pd.DataFrame(
                  y_pred,
                  columns=self._label_mapping[:n_labels],
                  index=adata.obs_names[indices],
              )
              return y_pred
  
      def train(
          self,
          max_epochs: Optional[int] = None,
          n_samples_per_label: Optional[float] = None,
          check_val_every_n_epoch: Optional[int] = None,
          train_size: float = 0.9,
          validation_size: Optional[float] = None,
          batch_size: int = 128,
          use_gpu: Optional[Union[str, int, bool]] = None,
          plan_kwargs: Optional[dict] = None,
          **trainer_kwargs,
      ):
          """
          Train the model.
  
          Parameters
          ----------
          max_epochs
              Number of passes through the dataset for semisupervised training.
          n_samples_per_label
              Number of subsamples for each label class to sample per epoch. By default, there
              is no label subsampling.
          check_val_every_n_epoch
              Frequency with which metrics are computed on the data for validation set for both
              the unsupervised and semisupervised trainers. If you'd like a different frequency for
              the semisupervised trainer, set check_val_every_n_epoch in semisupervised_train_kwargs.
          train_size
              Size of training set in the range [0.0, 1.0].
          validation_size
              Size of the test set. If `None`, defaults to 1 - `train_size`. If
              `train_size + validation_size < 1`, the remaining cells belong to a test set.
          batch_size
              Minibatch size to use during training.
          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), or use CPU (if False).
          plan_kwargs
              Keyword args for :class:`~scvi.train.SemiSupervisedTrainingPlan`. Keyword arguments passed to
              `train()` will overwrite values present in `plan_kwargs`, when appropriate.
          **trainer_kwargs
              Other keyword args for :class:`~scvi.train.Trainer`.
          """
  
          if max_epochs is None:
              n_cells = self.adata.n_obs
              max_epochs = np.min([round((20000 / n_cells) * 400), 400])
  
          if self.was_pretrained:
              max_epochs = int(np.min([10, np.max([2, round(max_epochs / 3.0)])]))
  
          logger.info("Training for {} epochs.".format(max_epochs))
  
          plan_kwargs = {} if plan_kwargs is None else plan_kwargs
  
          # if we have labeled cells, we want to subsample labels each epoch
          sampler_callback = (
              [SubSampleLabels()] if len(self._labeled_indices) != 0 else []
          )
  
          data_splitter = SemiSupervisedDataSplitter(
              adata=self.adata,
              unlabeled_category=self.unlabeled_category_,
              train_size=train_size,
              validation_size=validation_size,
              n_samples_per_label=n_samples_per_label,
              batch_size=batch_size,
              use_gpu=use_gpu,
          )
          training_plan = SemiSupervisedTrainingPlan(self.module, **plan_kwargs)
          if "callbacks" in trainer_kwargs.keys():
              trainer_kwargs["callbacks"].concatenate(sampler_callback)
          else:
              trainer_kwargs["callbacks"] = sampler_callback
  
          runner = TrainRunner(
              self,
              training_plan=training_plan,
              data_splitter=data_splitter,
              max_epochs=max_epochs,
              use_gpu=use_gpu,
              check_val_every_n_epoch=check_val_every_n_epoch,
              **trainer_kwargs,
          )
          return runner()
  