# Standard library imports
import logging
import time
import json
import pickle
from pathlib import Path
from typing import Optional, Union, Sequence, Mapping, Dict, List, Tuple, Any

# Third-party imports
import jax
import jax.numpy as jnp
from jax import jit, vmap, random, value_and_grad
from jax.scipy.stats import poisson
import optax
import tensorflow_probability.substrates.jax.distributions as tfd

import numpy as np
import pandas as pd
import scipy

import matplotlib
import seaborn as sns
from tqdm import tqdm

from anndata import AnnData
import anndata as ad

# Local imports
from scdef.utils import score_utils, color_utils
from scdef.utils.jax_utils import lognormal_sample, lognormal_entropy, gamma_logpdf

# Configure logging
logging.basicConfig()


class scDEF(object):
    """Single-cell Deep Exponential Families (scDEF) model.

    scDEF learns hierarchical, multi-level gene expression signatures from single-cell
    RNA-seq data provided in an AnnData object. This model can be used for a variety
    of analyses including dimensionality reduction, batch correction, clustering,
    and visualization of cell states and gene programs.

    The model fits multiple layers of latent factors ("gene signatures") to describe
    cellular heterogeneity at different resolutions. It supports batch correction,
    prior specification, and generation of corrected gene expression matrices.

    Model fitting, inference routines, and additional plotting utilities are
    implemented as methods of this class. The stored AnnData object is updated with
    model results during training.

    Args:
        adata: AnnData object containing the single-cell gene expression count matrix. Counts
            should be present in either `adata.X` or in the specified `adata.layers`.
        counts_layer: key for `adata.layers` specifying which layer to use as expression counts (if not `adata.X`).
        batch_key: key in `adata.obs` containing batch annotations; if provided, batch correction is performed.
            If None or not found, no batch correction is used.
        seed: random seed for model initialization and stochastic routines (uses JAX's pseudo-random number generator).
        n_factors: number of latent factors at the lowest layer (can be overridden by `layer_sizes`).
        decay_factor: size decay multiplier for the number of factors at each subsequent layer if `layer_sizes` not provided.
        max_n_layers: maximum number of hierarchical layers in the model.
        layer_sizes: explicit list of the number of factors in each scDEF layer. If None, layer sizes are set automatically.
        layer_names: list of custom names for the layers. If None, layer names are enumerated as ["L0", "L1", ...].
        logginglevel: verbosity level for the logger.
        alpha: concentration parameter of the Gamma distribution over the layer means
        shrinkage_shape: shape parameter for shrinkage prior controlling factor usage.
        shrinkage_rate: rate parameter for shrinkage prior controlling factor usage.
        top_alpha: concentration parameter for the top layer Dirichlet prior over factor proportions.
        factor_shape: shape of the prior distribution for factor-gene loadings matrix W.
        brd_strength: BRD (Batch Relevance Determination) prior concentration parameter for factor relevance estimation.
        brd_mean: mean of the BRD prior for factor relevance estimation.
        use_brd: if True, use BRD prior for automatic selection of active factors.
        cell_scale_shape: precision/concentration parameter for cell-specific scaling priors.
        gene_scale_shape: precision/concentration parameter for gene-specific scaling priors.
        batch_cpal: default matplotlib color palette name used for batches.
        layer_cpal: matplotlib color palette for factors/colors at each scDEF layer.
        lightness_mult: lightness multiplier to define the base color for each new scDEF layer.
        target_sum: target sum of the counts matrix.
        set_alpha_from_cov: if True, set the alpha parameter from the covariance of the counts matrix.
        marginalize_alpha: if True, marginalize the alpha parameter over the factors.
    """

    def make_corrected_data(self, layer_name: str = "scdef_corrected") -> None:
        """Compute and store the low-rank reconstruction of the UMI count matrix.

        The reconstructed matrix is saved to adata.layers[layer_name], providing a
        denoised, batch-corrected version of the expression data.

        Args:
            layer_name: name for the AnnData layer where the reconstructed matrix is stored
        """
        scdef_layer = self.layer_names[0]
        Z = self.pmeans[f"{scdef_layer}z"][:, self.factor_lists[0]]  # nxk
        W = self.pmeans[f"{scdef_layer}W"][self.factor_lists[0]]  # kxg
        self.adata.layers[layer_name] = np.array(Z.dot(W))

    def __init__(
        self,
        adata: AnnData,
        counts_layer: Optional[str] = None,
        batch_key: Optional[str] = None,
        seed: Optional[int] = 42,
        n_factors: Optional[int] = 100,
        decay_factor: Optional[float] = 2.0,
        max_n_layers: Optional[float] = 5,
        layer_sizes: Optional[list] = None,
        layer_names: Optional[list] = None,
        logginglevel: Optional[int] = logging.INFO,
        alpha: Optional[float] = 1.0,
        shrinkage_shape: Optional[float] = 1.0,
        shrinkage_rate: Optional[float] = 1.0,
        shrinkage_mean: Optional[float] = 1.0,
        top_alpha: Optional[float] = 1.0,
        factor_shape: Optional[float] = 0.1,
        brd_strength: Optional[float] = 1.0,
        brd_mean: Optional[float] = 1.0,
        use_brd: Optional[bool] = True,
        cell_scale_shape: Optional[float] = 1.0,
        gene_scale_shape: Optional[float] = 1.0,
        batch_cpal: Optional[str] = "Dark2",
        layer_cpal: Optional[str] = "tab10",
        lightness_mult: Optional[float] = 0.15,
        target_sum: Optional[float] = 1_000,
        set_alpha_from_cov: Optional[bool] = False,
        marginalize_alpha: Optional[bool] = False,
    ):
        self.n_cells, self.n_genes = adata.shape

        self.n_batches = 1
        self.batches = [""]
        self.batch_key = batch_key

        self.seed = seed
        self.batch_cpal = batch_cpal
        self.layer_cpal = layer_cpal
        self.lightness_mult = lightness_mult

        self.logger = logging.getLogger("scDEF")
        self.logger.setLevel(logginglevel)

        self.load_adata(adata, layer=counts_layer, batch_key=batch_key)

        self.top_alpha = top_alpha
        self.alpha = alpha
        self.factor_shape = factor_shape
        self.brd = brd_strength
        self.brd_mean = brd_mean
        self.shrinkage_shape = shrinkage_shape
        self.shrinkage_rate = shrinkage_rate
        self.shrinkage_mean = shrinkage_mean
        self.cell_scale_shape = cell_scale_shape
        self.gene_scale_shape = gene_scale_shape
        self.use_brd = use_brd
        self.target_sum = target_sum
        self.set_alpha_from_cov = set_alpha_from_cov
        self.marginalize_alpha = marginalize_alpha

        self.decay_factor = decay_factor
        self.n_factors = n_factors
        self.max_n_layers = max_n_layers
        if layer_sizes is not None:
            self.layer_sizes = [int(x) for x in layer_sizes]
            self.n_factors = int(layer_sizes[0])
            self.n_layers = len(self.layer_sizes)
        else:
            self.update_model_size(n_factors)

        if layer_names is not None:
            self.layer_names = layer_names
        else:
            self.layer_names = [f"L{i}" for i in range(self.n_layers)]
        self.factor_lists = [np.arange(size) for size in self.layer_sizes]
        self.set_factor_names()

        self.update_model_priors()

        self.make_layercolors(layer_cpal=self.layer_cpal, lightness_mult=lightness_mult)

        self.init_var_params(nmf_init=False)  # just to get stub
        self.set_posterior_means()
        self.set_posterior_variances()
        self._has_fit = False
        self._fit_revision = 0

    def __repr__(self):
        out = f"scDEF object with {self.n_layers} layers"
        out += (
            "\n\t"
            + "Layer names: "
            + ", ".join([f"{name}" for name in self.layer_names])
        )
        out += (
            "\n\t"
            + "Layer sizes: "
            + ", ".join([str(len(factors)) for factors in self.factor_lists])
        )
        out += "\n\t" + "Alpha parameter: " + str(self.alpha)
        # out += (
        #     "\n\t"
        #     + "Layer factor shape parameters: "
        #     + ", ".join([str(shape) for shape in self.factor_shapes])
        # )
        if self.use_brd:
            out += "\n\t" + "Using BRD"
        out += "\n\t" + "Number of batches: " + str(self.n_batches)
        out += "\n" + "Contains " + self.adata.__str__()
        return out

    def save(
        self,
        dir_path: Union[str, Path],
        overwrite: bool = False,
        save_anndata: bool = False,
    ) -> None:
        """Save model state to disk, similarly to scvi-tools.

        This writes a model state pickle plus metadata to ``dir_path``. AnnData
        is saved separately as ``adata.h5ad`` only when ``save_anndata=True``.

        Args:
            dir_path: output directory path
            overwrite: whether to overwrite an existing non-empty directory
            save_anndata: whether to save ``model.adata`` as ``adata.h5ad``
        """
        out_dir = Path(dir_path).expanduser().resolve()
        if out_dir.exists() and any(out_dir.iterdir()) and not overwrite:
            raise FileExistsError(
                f"Directory '{out_dir}' already exists and is not empty. "
                "Set overwrite=True to overwrite."
            )
        out_dir.mkdir(parents=True, exist_ok=True)

        if save_anndata:
            self.adata.write(out_dir / "adata.h5ad")

        state = dict(self.__dict__)
        # Graphviz graph objects are environment-dependent and not required for
        # restoring model state.
        if "graph" in state:
            state["graph"] = None
        logger_level = (
            self.logger.level
            if hasattr(self, "logger") and self.logger is not None
            else None
        )
        state["logger"] = None

        payload = {
            "state": state,
            "class_name": self.__class__.__name__,
            "module": self.__class__.__module__,
            "save_anndata": bool(save_anndata),
            "logger_level": logger_level,
        }
        with open(out_dir / "model_state.pkl", "wb") as f:
            pickle.dump(payload, f, protocol=pickle.HIGHEST_PROTOCOL)

        metadata = {
            "class_name": payload["class_name"],
            "module": payload["module"],
            "save_anndata": payload["save_anndata"],
            "files": ["model_state.pkl"] + (["adata.h5ad"] if save_anndata else []),
        }
        with open(out_dir / "metadata.json", "w", encoding="utf-8") as f:
            json.dump(metadata, f, indent=2)

    @classmethod
    def load(
        cls,
        dir_path: Union[str, Path],
        adata: Optional[AnnData] = None,
    ) -> "scDEF":
        """Load model from disk.

        Args:
            dir_path: directory created by :meth:`save`
            adata: optional AnnData to attach when ``adata.h5ad`` was not saved.

        Returns:
            Loaded model instance.
        """
        in_dir = Path(dir_path).expanduser().resolve()
        state_path = in_dir / "model_state.pkl"
        if not state_path.exists():
            raise FileNotFoundError(f"Model state file not found: '{state_path}'.")

        with open(state_path, "rb") as f:
            try:
                payload = pickle.load(f)
            except ModuleNotFoundError as e:
                # Compatibility path for artifacts pickled under a different
                # NumPy internal module layout (e.g. numpy._core vs numpy.core).
                if "numpy._core" not in str(e):
                    raise
                f.seek(0)

                class _GraphPlaceholder:
                    """Fallback placeholder for optional Graphviz graph classes."""

                    def __init__(self, *args, **kwargs):
                        self.args = args
                        self.kwargs = kwargs

                class _CompatUnpickler(pickle.Unpickler):
                    def find_class(self, module, name):
                        if module.startswith("numpy._core"):
                            module = module.replace("numpy._core", "numpy.core", 1)
                        if module == "numpy.rec":
                            module = "numpy.core.records"
                        if module == "graphviz.graphs" and name in {"Graph", "Digraph"}:
                            return _GraphPlaceholder
                        return super().find_class(module, name)

                payload = _CompatUnpickler(f).load()

        obj = cls.__new__(cls)
        obj.__dict__.update(payload["state"])
        obj.logger = logging.getLogger(payload.get("class_name", cls.__name__))
        if payload.get("logger_level") is not None:
            obj.logger.setLevel(payload["logger_level"])

        adata_path = in_dir / "adata.h5ad"
        if adata is not None:
            obj.adata = adata
        elif getattr(obj, "adata", None) is not None:
            # Prefer embedded AnnData saved in the model state.
            pass
        elif payload.get("save_anndata", False) and adata_path.exists():
            obj.adata = ad.read_h5ad(adata_path)
        else:
            raise ValueError(
                "No AnnData available. Provide `adata` to load(...), or save with "
                "`save_anndata=True` so `adata.h5ad` is available."
            )

        obj.n_cells, obj.n_genes = obj.adata.shape
        return obj

    def make_layercolors(
        self,
        layer_cpal: Optional[str] = "tab10",
        lightness_mult: Optional[float] = 0.15,
    ):
        if isinstance(layer_cpal, str):
            layer_cpal = [layer_cpal] * self.n_layers
        elif isinstance(layer_cpal, list):
            if len(layer_cpal) != self.n_layers:
                raise ValueError("layer_cpal list must be of size scDEF.n_layers")
        else:
            raise TypeError("layer_cpal must be either a cmap str or a list of strs")

        layer_sizes = []
        for i in range(self.n_layers):
            layer_size = len(self.factor_lists[i])
            if layer_size > 10:
                layer_cpal[i] = "tab20"
            elif layer_size > 20:
                layer_cpal[i] = "hsl"
            layer_sizes.append(layer_size)

        self.layer_colorpalettes = [
            sns.color_palette(layer_cpal[idx], n_colors=size)
            for idx, size in enumerate(layer_sizes)
        ]

        if len(np.unique(layer_cpal)) == 1:
            # Make the layers have different lightness
            for layer_idx, size in enumerate(layer_sizes):
                for factor_idx in range(size):
                    col = self.layer_colorpalettes[layer_idx][factor_idx]
                    self.layer_colorpalettes[layer_idx][
                        factor_idx
                    ] = color_utils.adjust_lightness(
                        col, amount=1.0 + lightness_mult * layer_idx
                    )

    def load_adata(self, adata, layer=None, batch_key=None):
        if not isinstance(adata, AnnData):
            raise TypeError("adata must be an instance of AnnData.")
        self.adata = adata.copy()
        self.adata.raw = self.adata
        X = self.adata.raw.X
        if layer is not None:
            X = self.adata.layers[layer]

        if isinstance(X, scipy.sparse.csr_matrix):
            self.X = np.array(X.toarray())
        else:
            self.X = np.array(X)

        self.X = self.X.astype(float)

        self.batch_indices_onehot = np.ones((self.adata.shape[0], 1))
        self.batch_lib_sizes = np.sum(self.X, axis=1)
        self.batch_lib_ratio = (
            np.ones((self.X.shape[0], 1))
            * np.mean(self.batch_lib_sizes)
            / np.var(self.batch_lib_sizes)
        )
        self.exposure = jnp.array(self.batch_lib_sizes / np.mean(self.batch_lib_sizes))

        eps = 1e-12

        # ---- NEW: gene_ratio is per-gene and encodes baseline abundance ----
        # Choose a baseline definition. For raw UMI counts, using mean counts per cell is reasonable.
        # If you prefer library-normalized baseline, replace mu = X.mean(0) by mu = (X/lib).mean(0).
        mu = self.X.mean(axis=0) + eps
        mu_bar = float(mu.mean())
        # Prior mean of gene_scale_g will be ~ mu_g / mu_bar
        gene_ratio = mu_bar / mu  # vector length G
        gene_ratio = np.array(gene_ratio)[None, :]  # shape (1, G)
        self.gene_ratio_init = (
            gene_ratio  # will be overridden per batch if batch_key has >1 values
        )

        gene_size = np.sum(self.X, axis=0)
        self.gene_means = np.mean(gene_size)
        self.gene_vars = np.var(gene_size)
        self.gene_ratio = self.gene_means / self.gene_vars

        if batch_key is not None:
            if batch_key in self.adata.obs.columns:
                batches = np.unique(self.adata.obs[batch_key].values)
                self.batches = batches
                self.n_batches = len(batches)
                self.logger.info(
                    f"Found {self.n_batches} values for `{batch_key}` in data: {batches}"
                )
                if f"{batch_key}_colors" in self.adata.uns:
                    self.batch_colors = self.adata.uns[f"{batch_key}_colors"]
                else:
                    batch_colorpalette = sns.color_palette(
                        self.batch_cpal, n_colors=self.n_batches
                    )
                    self.batch_colors = [
                        matplotlib.colors.to_hex(batch_colorpalette[idx])
                        for idx in range(self.n_batches)
                    ]
                    self.adata.uns[f"{batch_key}_colors"] = self.batch_colors

                self.batch_indices_onehot = np.zeros(
                    (self.adata.shape[0], self.n_batches)
                )
                # If multiple batches: compute gene_ratio per batch (still per gene)
                if self.n_batches > 1:
                    self.gene_ratio = np.ones((self.n_batches, self.adata.shape[1]))
                    self.gene_ratio_init = np.ones(
                        (self.n_batches, self.adata.shape[1])
                    )
                    for i, b in enumerate(batches):
                        cells = np.where(self.adata.obs[batch_key] == b)[0]
                        self.batch_indices_onehot[cells, i] = 1
                        self.batch_lib_sizes[cells] = np.sum(self.X, axis=1)[cells]
                        self.batch_lib_ratio[cells] = np.mean(
                            self.batch_lib_sizes[cells]
                        ) / (np.var(self.batch_lib_sizes[cells]) + eps)
                        mu_b = self.X[cells].mean(axis=0) + eps
                        mu_b_bar = float(mu_b.mean())
                        self.gene_ratio_init[i] = mu_b_bar / mu_b
                        gene_size = np.sum(self.X[cells], axis=0)
                        gene_means = np.mean(gene_size)
                        gene_vars = np.var(gene_size)
                        self.gene_ratio[i] = (
                            self.gene_ratio[i] * 0 + gene_means / gene_vars
                        )

        self.batch_indices_onehot = jnp.array(self.batch_indices_onehot)
        self.batch_lib_sizes = jnp.array(self.batch_lib_sizes)
        self.batch_lib_ratio = jnp.array(self.batch_lib_ratio)
        self.gene_ratio = jnp.array(self.gene_ratio)

    def update_model_size(
        self,
        max_n_factors,
        max_n_layers=None,
        layer_sizes=None,
    ):
        """Update latent hierarchy dimensions.

        Args:
            max_n_factors: maximum number of factors when ``layer_sizes`` is not
                provided.
            max_n_layers: maximum number of layers when ``layer_sizes`` is not
                provided.
            layer_sizes: explicit per-layer sizes. If provided, sizes are
                sanitized to be non-increasing and consecutive duplicates are
                collapsed.
        """
        if layer_sizes is not None:
            # If two consecutive layers have increasing sizes, clip
            for i in range(len(layer_sizes) - 1):
                if layer_sizes[i + 1] > layer_sizes[i]:
                    layer_sizes[i + 1] = layer_sizes[i]
            # If consecutive layers have the same size, keep only one.
            # Build a deduplicated list directly so chains like [...,1,1,1]
            # collapse to a single 1.
            dedup_layer_sizes = []
            for size in layer_sizes:
                if len(dedup_layer_sizes) == 0 or size != dedup_layer_sizes[-1]:
                    dedup_layer_sizes.append(size)
            layer_sizes = dedup_layer_sizes
            self.layer_sizes = layer_sizes
            self.n_layers = len(self.layer_sizes)
            self.layer_names = [f"L{i}" for i in range(self.n_layers)]
            self.factor_lists = [np.arange(size) for size in self.layer_sizes]
            self.set_factor_names()
            return

        if max_n_layers is None:
            max_n_layers = self.max_n_layers
        n_factors = int(max_n_factors)
        self.layer_sizes = []
        self.layer_sizes.append(n_factors)
        if max_n_layers > 1:
            while len(self.layer_sizes) < max_n_layers:
                n_factors = int(np.floor(n_factors / self.decay_factor))
                self.layer_sizes.append(n_factors)
                if n_factors == 1:
                    break
            if self.layer_sizes[-1] > 1:
                self.layer_sizes[-1] = 1

        self.n_layers = len(self.layer_sizes)
        self.layer_names = [f"L{i}" for i in range(self.n_layers)]
        self.factor_lists = [np.arange(size) for size in self.layer_sizes]
        self.set_factor_names()

    def update_model_priors(self):
        if self.use_brd:
            self.factor_shapes = [1.0] + [
                self.factor_shape * self.layer_sizes[layer_idx] / self.layer_sizes[0]
                for layer_idx in range(1, self.n_layers)
            ]
        else:
            self.factor_shapes = [
                self.factor_shape * self.layer_sizes[layer_idx] / self.layer_sizes[0]
                for layer_idx in range(self.n_layers)
            ]

        self.w_priors = []
        for idx in range(self.n_layers):
            prior_shapes = self.factor_shapes[idx]
            prior_rates = self.factor_shapes[idx]

            if idx > 0:
                prior_shapes = (
                    np.ones((self.layer_sizes[idx], self.layer_sizes[idx - 1]))
                    * self.factor_shapes[idx]
                    * 1.0
                )
                prior_rates = (
                    np.ones((self.layer_sizes[idx], self.layer_sizes[idx - 1]))
                    * self.factor_shapes[idx]
                )
                prior_shapes = jnp.clip(jnp.array(prior_shapes), 1e-12, 1e12)
                prior_rates = jnp.clip(jnp.array(prior_rates), 1e-12, 1e12)

            self.w_priors.append([prior_shapes, prior_rates])

        if self.set_alpha_from_cov:
            self.alpha = self.alpha * self.batch_lib_sizes.mean() / self.target_sum

    def get_effective_factors(
        self,
        brd_min: Optional[float] = 1.0,
        ard_min: Optional[float] = 0.001,
        clarity_min: Optional[float] = 0.5,
        min_cells: Optional[float] = 0.001,
    ):
        layer_name = self.layer_names[0]
        if min_cells != 0:
            min_cells = (
                np.maximum(10, min_cells * self.n_cells)
                if min_cells < 1.0
                else min_cells
            )

        normed = (
            self.pmeans[f"{layer_name}z"]
            / np.sum(self.pmeans[f"{layer_name}z"], axis=1)[:, None]
        )
        assignments = np.argmax(normed, axis=1)
        counts = np.array(
            [np.count_nonzero(assignments == a) for a in range(self.layer_sizes[0])]
        )

        keep = np.array(range(self.layer_sizes[0]))[np.where(counts >= min_cells)[0]]
        if self.use_brd:
            required_cols = {"BRD", "ARD", "clarity_score_01", "original_factor_idx"}
            if "factor_obs" not in self.adata.uns or not required_cols.issubset(
                set(self.adata.uns["factor_obs"].columns)
            ):
                from scdef.tools.factor import factor_diagnostics

                factor_diagnostics(self)

            factor_obs = self.adata.uns["factor_obs"]
            if "child_layer" in factor_obs.columns:
                factor_obs_l0 = factor_obs[
                    factor_obs["child_layer"] == self.layer_names[0]
                ]
            else:
                factor_obs_l0 = factor_obs

            brd = np.full(self.layer_sizes[0], np.nan, dtype=float)
            ard = np.full(self.layer_sizes[0], np.nan, dtype=float)
            clarity = np.full(self.layer_sizes[0], np.nan, dtype=float)

            original_idx = factor_obs_l0["original_factor_idx"].to_numpy(dtype=int)
            brd_vals = factor_obs_l0["BRD"].to_numpy(dtype=float)
            ard_vals = factor_obs_l0["ARD"].to_numpy(dtype=float)
            clarity_vals = factor_obs_l0["clarity_score_01"].to_numpy(dtype=float)

            valid_idx = (original_idx >= 0) & (original_idx < self.layer_sizes[0])
            original_idx = original_idx[valid_idx]
            brd_vals = brd_vals[valid_idx]
            ard_vals = ard_vals[valid_idx]
            clarity_vals = clarity_vals[valid_idx]

            brd[original_idx] = brd_vals
            ard[original_idx] = ard_vals
            clarity[original_idx] = clarity_vals

            valid = np.isfinite(brd) & np.isfinite(ard) & np.isfinite(clarity)
            ard_sum = np.nansum(ard)
            brd_keep = np.where(
                valid
                & (brd >= brd_min)
                & (clarity >= clarity_min)
                & (ard >= ard_min * ard_sum)
            )[0]
            keep = np.unique(list(set(brd_keep).intersection(keep)))

        return keep

    def init_var_params(
        self,
        init_budgets=True,
        init_alpha=True,
        init_z=None,
        init_w=None,
        init_brd=None,
        init_ard=None,
        nmf_init=False,
        z_init_concentration=1.0,
        **kwargs,
    ):
        rngs = random.split(random.PRNGKey(self.seed), self.n_layers)

        def _get_layer_init(init_value, layer_idx):
            if init_value is None:
                return None
            if isinstance(init_value, (list, tuple)):
                if layer_idx >= len(init_value) or init_value[layer_idx] is None:
                    return None
                return np.asarray(init_value[layer_idx], dtype=np.float32)
            if layer_idx == 0:
                return np.asarray(init_value, dtype=np.float32)
            return None

        if init_budgets:
            m = np.array(self.batch_lib_sizes / np.mean(self.batch_lib_sizes))[
                :, None
            ]  # (N,1)
            m = np.clip(m, 1e-3, 1e2)
            v = m / 10.0
            self.local_params = [
                jnp.array(
                    (
                        jnp.log(m**2 / jnp.sqrt(m**2 + v))
                        * jnp.ones((self.n_cells, 1)),  # cell_scales
                        jnp.log(jnp.sqrt(jnp.log(1 + v / (m**2))))
                        * jnp.ones((self.n_cells, 1)),
                    )
                ),
            ]
            # initialize gene scales at the prior mean (per gene)
            # prior mean under Gamma(shape, rate=shape*gene_ratio) is 1/gene_ratio
            m = 1.0 / np.array(self.gene_ratio_init)  # shape: (G,) or (B,G)
            m = np.clip(m, 1e-6, 1e6)
            v = m / 10.0

            m = jnp.array(m, dtype=jnp.float32)
            v = jnp.array(v, dtype=jnp.float32)

            self.global_params = [
                jnp.array(
                    (
                        jnp.log(m**2 / jnp.sqrt(m**2 + v)),
                        jnp.log(jnp.sqrt(jnp.log(1 + v / (m**2)))),
                    )
                ),
            ]
        else:
            self.local_params = [self.local_params[0]]
            self.global_params = [self.global_params[0]]
        # BRD
        m_brd = tfd.Gamma(100.0, 100.0 / self.brd_mean).sample(
            seed=rngs[0],
            sample_shape=[self.layer_sizes[0], 1],
        )  # self.brd_mean
        v_brd = m_brd / 100.0
        if init_brd is not None:
            m_brd = init_brd
            v_brd = m_brd / 10.0
        self.global_params.append(
            jnp.array(
                (
                    jnp.log(m_brd**2 / jnp.sqrt(m_brd**2 + v_brd))
                    * jnp.ones((self.layer_sizes[0], 1)),
                    jnp.log(jnp.sqrt(jnp.log(1 + v_brd / (m_brd**2))))
                    * jnp.ones((self.layer_sizes[0], 1)),
                )
            ),
        )

        if nmf_init and init_w is None:
            self.logger.info("Initializing the factor weights with NMF.")
            init_z, init_w = self.get_nmf_init(**kwargs)

        z_shapes = []
        z_rates = []
        rng_cnt = 0

        for layer_idx in range(
            self.n_layers
        ):  # we go from layer 0 (bottom) to layer L (top)
            # Init z
            a = z_init_concentration
            clip = 1e-6  # alpha=1 and clip=1e-6 allows for many factors to be learned, but leads to BRD becoming kind of big

            if (
                layer_idx > 0
            ):  # If the top layer inits to very small numbers, the loss goes crazy when MC=10...
                clip = 1e-6
                a = 1.0

            if layer_idx == self.n_layers - 1:
                clip = 1e-3

            z_init_layer = _get_layer_init(init_z, layer_idx)
            if z_init_layer is not None and not nmf_init:
                m = z_init_layer.astype(jnp.float32)
            else:
                m = jnp.clip(
                    tfd.Gamma(a, a / 1.0).sample(
                        seed=rngs[rng_cnt],
                        sample_shape=[self.n_cells, self.layer_sizes[layer_idx]],
                    ),
                    clip,
                    1e1,
                )
                rng_cnt += 1

            v = m  # / 100.0
            if layer_idx > 0:
                v = m  # / 100.0
            z_shape = jnp.log(m**2 / jnp.sqrt(m**2 + v)) * jnp.ones(
                (self.n_cells, self.layer_sizes[layer_idx])
            )
            z_shapes.append(z_shape)
            z_rate = jnp.log(jnp.sqrt(jnp.log(1 + v / (m**2)))) * jnp.ones(
                (self.n_cells, self.layer_sizes[layer_idx])
            )
            z_rates.append(z_rate)

        self.local_params.append(jnp.array((jnp.hstack(z_shapes), jnp.hstack(z_rates))))

        for layer_idx in range(
            self.n_layers
        ):  # the w don't have a shared axis across all, so we can't vectorize
            # Init w
            in_layer = self.layer_sizes[layer_idx]
            if layer_idx == 0:
                out_layer = self.n_genes
            else:
                out_layer = self.layer_sizes[layer_idx - 1]

            a = 100.0

            w_init_layer = _get_layer_init(init_w, layer_idx)
            if w_init_layer is not None and not nmf_init:
                m = w_init_layer.astype(jnp.float32)
                v = m
            else:
                m = 1.0 / self.layer_sizes[layer_idx] * jnp.ones((in_layer, out_layer))
                m = m * self.w_priors[layer_idx][0] / self.w_priors[layer_idx][1]
                m = tfd.Gamma(100.0, 100.0 / (m)).sample(seed=rngs[rng_cnt])
                rng_cnt += 1
                v = m / 100.0
            if nmf_init and layer_idx == 0:
                iw = init_w[layer_idx].astype(jnp.float32)
                # Rescale
                iw = 1.0 * iw + 1.0
                # iw = 10. * 100.**iw/100. + 1.
                iw = iw
                iw = iw  # * self.w_priors[layer_idx][0] / self.w_priors[layer_idx][1]
                iw = (
                    iw / np.sum(iw, axis=1, keepdims=True)
                    + 1.0 / self.layer_sizes[layer_idx]
                )
                iw = jnp.clip(iw, 1e-6, 1e2)  # need to set scales to avoid huge BRDs...
                m = tfd.Gamma(100.0, 100.0 / (iw)).sample(seed=rngs[rng_cnt])
                rng_cnt += 1
                v = m / 100.0

            w_shape = jnp.log(m**2 / jnp.sqrt(m**2 + v)) * jnp.ones(
                (in_layer, out_layer)
            )
            w_rate = jnp.log(jnp.sqrt(jnp.log(1 + v / (m**2)))) * jnp.ones(
                (in_layer, out_layer)
            )

            self.global_params.append(jnp.array((w_shape, w_rate)))

        # Gamma(1e3, 1e3/m)
        m = (
            self.shrinkage_shape
        )  # self.shrinkage_mean * self.shrinkage_shape  / self.shrinkage_rate
        v = m / 10.0
        if init_ard is not None:
            m = init_ard
            v = m / 10.0

        self.global_params.append(
            jnp.array(
                (
                    jnp.log(m**2 / jnp.sqrt(m**2 + v))
                    * jnp.ones((self.layer_sizes[0], 1)),
                    jnp.log(jnp.sqrt(jnp.log(1 + v / (m**2))))
                    * jnp.ones((self.layer_sizes[0], 1)),
                )
            ),
        )

        if init_alpha:
            m = self.alpha  # + 0. * np.array(self.gene_ratio_init)
        else:
            m = self.pmeans["alpha"]  # + 0. * np.array(self.gene_ratio_init)
        v = m / 10.0

        self.global_params.append(
            jnp.array(
                (
                    jnp.log(m**2 / jnp.sqrt(m**2 + v)) * jnp.ones((1, 1)),
                    jnp.log(jnp.sqrt(jnp.log(1 + v / (m**2)))) * jnp.ones((1, 1)),
                )
            ),
        )

    def get_nmf_init(self, max_cells=None):
        """Use NMF on the data to initialize the first layer and then recursively for the other layers.

        Args:
            max_cells: maximum number of cells to use for NMF initialization

        Returns:
            tuple of (init_z, init_W) initialization values
        """
        init_z = []
        init_W = []
        from sklearn.decomposition import NMF

        X = self.X
        if max_cells is not None:
            if max_cells < X.shape[0]:
                key = jax.random.PRNGKey(self.seed)
                cells = jax.random.choice(
                    key, X.shape[0], replace=False, shape=(max_cells,)
                )
                X = X[cells]
        X = X.copy()
        lib = np.sum(X, axis=1, keepdims=True)
        X /= lib
        X *= 10_000
        # X = np.log(X + 1)
        for layer_idx in range(self.n_layers):
            model = NMF(
                n_components=self.layer_sizes[layer_idx],
                random_state=self.seed,
                beta_loss="kullback-leibler",
                solver="mu",
            )
            z = model.fit_transform(X)
            W = model.components_
            X = z
            init_z.append(z)
            init_W.append(W)

        return init_z, init_W

    def identify_mixture_factors(
        self, max_n_genes: int = 20, thres: float = 0.5
    ) -> np.ndarray:
        """Identify factors that might be better if broken apart.

        Args:
            max_n_genes: maximum number of genes per factor
            thres: threshold for identifying mixture factors

        Returns:
            array of factor indices that are mixture factors
        """
        # sparse_factors = np.where(self.pmeans['brd'].ravel() > 1.)[0]
        kept_factors = self.factor_lists[0]
        normed_factors = (
            self.pmeans["L0W"][kept_factors]
            / self.pmeans["L0W"][kept_factors].max(axis=1)[:, None]
        )
        n_genes_per_factor = np.sum(normed_factors > thres, axis=1)
        mixture_factors = kept_factors[np.where(n_genes_per_factor > max_n_genes)[0]]
        return mixture_factors

    def reinit_factors(
        self, mixture_factors=None, init_budgets=False, exponent=1.1, **kwargs
    ):
        kept_factors = self.factor_lists[0]
        if mixture_factors is None:
            mixture_factors = self.identify_mixture_factors(**kwargs)
        # Make W0 matrix with sparse factors, and with each mixture factor repeated twice
        W0 = np.ones((self.layer_sizes[0], self.n_genes)) * 1.0 / self.layer_sizes[0]
        for i, factor in enumerate(kept_factors):
            W0[i] = self.pmeans["L0W"][factor] ** exponent

        # If we can fit repetitions
        if len(mixture_factors) <= self.layer_sizes[0] - len(kept_factors):
            for factor in mixture_factors:
                # Add another copy
                W0[len(kept_factors) + i] = self.pmeans["L0W"][factor] ** exponent

        self.init_var_params(init_budgets=init_budgets, init_w=W0, nmf_init=False)
        return W0

    def elbo(
        self,
        rng,
        batch,
        indices,
        local_params,
        global_params,
        annealing_parameter,
        stop_gradients,
        stop_cell_budgets,
        stop_gene_budgets,
        min_shape=jnp.log(1e-10),
        max_shape=jnp.log(1e6),
        min_rate=jnp.log(1e-10),
        max_rate=jnp.log(1e10),
    ):
        # Single-sample Monte Carlo estimate of the variational lower bound.
        batch_indices_onehot = self.batch_indices_onehot[indices]

        cell_budget_params = local_params[0]
        z_params = local_params[1]
        gene_budget_params = global_params[0]
        fscale_params = global_params[1]

        cell_budget_shape = jnp.clip(
            cell_budget_params[0][indices], min_shape, max_shape
        )
        cell_budget_rate = jnp.exp(
            jnp.clip(cell_budget_params[1][indices], min_rate, max_rate)
        )
        cell_budget_shape = jax.lax.cond(
            stop_gradients[0],
            lambda: jax.lax.stop_gradient(cell_budget_shape),
            lambda: cell_budget_shape,
        )
        cell_budget_rate = jax.lax.cond(
            stop_gradients[0],
            lambda: jax.lax.stop_gradient(cell_budget_rate),
            lambda: cell_budget_rate,
        )

        gene_budget_shape = jnp.maximum(gene_budget_params[0], min_shape)
        gene_budget_rate = jnp.exp(jnp.clip(gene_budget_params[1], min_rate, max_rate))
        gene_budget_shape = jax.lax.cond(
            stop_gradients[0],
            lambda: jax.lax.stop_gradient(gene_budget_shape),
            lambda: gene_budget_shape,
        )
        gene_budget_rate = jax.lax.cond(
            stop_gradients[0],
            lambda: jax.lax.stop_gradient(gene_budget_rate),
            lambda: gene_budget_rate,
        )

        fscale_shapes = jnp.clip(fscale_params[0], min_shape, max_shape)
        fscale_rates = jnp.exp(jnp.clip(fscale_params[1], min_rate, max_rate))
        fscale_shapes = jax.lax.cond(
            stop_gradients[0],
            lambda: jax.lax.stop_gradient(fscale_shapes),
            lambda: fscale_shapes,
        )
        fscale_rates = jax.lax.cond(
            stop_gradients[0],
            lambda: jax.lax.stop_gradient(fscale_rates),
            lambda: fscale_rates,
        )

        wm_shape = jnp.clip(global_params[2 + self.n_layers][0], min_shape, max_shape)
        wm_rate = jnp.exp(
            jnp.clip(global_params[2 + self.n_layers][1], min_rate, max_rate)
        )

        s_shape = jnp.clip(
            global_params[2 + self.n_layers + 1][0], min_shape, max_shape
        )
        s_rate = jnp.exp(
            jnp.clip(global_params[2 + self.n_layers + 1][1], min_rate, max_rate)
        )

        z_shapes = jnp.clip(z_params[0][indices], min_shape, max_shape)
        z_rates = jnp.exp(jnp.clip(z_params[1][indices], min_rate, max_rate))

        # Sample from variational distribution
        cell_budget_shape = jax.lax.cond(
            stop_cell_budgets,
            lambda: jax.lax.stop_gradient(cell_budget_shape),
            lambda: cell_budget_shape,
        )
        cell_budget_rate = jax.lax.cond(
            stop_cell_budgets,
            lambda: jax.lax.stop_gradient(cell_budget_rate),
            lambda: cell_budget_rate,
        )
        cell_budget_sample = lognormal_sample(rng, cell_budget_shape, cell_budget_rate)
        gene_budget_shape = jax.lax.cond(
            stop_gene_budgets,
            lambda: jax.lax.stop_gradient(gene_budget_shape),
            lambda: gene_budget_shape,
        )
        gene_budget_rate = jax.lax.cond(
            stop_gene_budgets,
            lambda: jax.lax.stop_gradient(gene_budget_rate),
            lambda: gene_budget_rate,
        )
        gene_budget_sample = lognormal_sample(rng, gene_budget_shape, gene_budget_rate)
        if self.use_brd:
            fscale_samples = lognormal_sample(rng, fscale_shapes, fscale_rates)
            _wm_sample = lognormal_sample(rng, wm_shape, wm_rate)
            brd_samples = fscale_samples  # / jnp.maximum(_wm_sample, 1.0)
            # log_brd = jnp.log(fscale_samples + 1e-8)
            # log_brd = log_brd - jnp.mean(log_brd) + jnp.log(self.brd_mean)
            # brd_samples = jnp.exp(log_brd)
        # w will be sampled in a loop below because it cannot be vectorized

        # Compute ELBO
        global_pl = gamma_logpdf(
            gene_budget_sample,
            self.gene_scale_shape,
            self.gene_scale_shape * self.gene_ratio,
        )
        global_en = lognormal_entropy(gene_budget_shape, gene_budget_rate)
        # exposure = (self.batch_lib_sizes)[:, None]
        # local_pl = 0.0
        # local_en = 0.0
        # local_pl = gamma_logpdf(cell_budget_sample,
        #                         self.cell_scale_shape,
        #                         self.cell_scale_shape / exposure[indices])
        local_pl = gamma_logpdf(
            cell_budget_sample,
            self.cell_scale_shape,
            self.cell_scale_shape * self.batch_lib_ratio[indices],
        )
        local_en = lognormal_entropy(cell_budget_shape, cell_budget_rate)

        # scale
        s_sample = lognormal_sample(rng, s_shape, s_rate)
        global_pl += gamma_logpdf(s_sample, 100.0, 100.0 / self.alpha)
        global_en += lognormal_entropy(s_shape, s_rate)
        if self.use_brd:
            global_pl += gamma_logpdf(
                fscale_samples, self.brd, self.brd / self.brd_mean
            )
            global_en += lognormal_entropy(fscale_shapes, fscale_rates)

            global_pl += gamma_logpdf(
                _wm_sample, self.shrinkage_shape, self.shrinkage_rate
            )

            global_en += lognormal_entropy(wm_shape, wm_rate)

        z_mean = 1.0
        for idx in list(np.arange(0, self.n_layers)[::-1]):
            start = np.sum(self.layer_sizes[:idx]).astype(int)
            end = start + self.layer_sizes[idx]

            # w
            _w_shape = jnp.clip(global_params[2 + idx][0], min_shape, max_shape)
            _w_rate = jnp.exp(jnp.clip(global_params[2 + idx][1], min_rate, max_rate))

            _w_shape = jax.lax.cond(
                stop_gradients[idx],
                lambda: jax.lax.stop_gradient(_w_shape),
                lambda: _w_shape,
            )
            _w_rate = jax.lax.cond(
                stop_gradients[idx],
                lambda: jax.lax.stop_gradient(_w_rate),
                lambda: _w_rate,
            )

            _w_sample = lognormal_sample(rng, _w_shape, _w_rate)

            if idx == 0 and self.use_brd:
                global_pl += gamma_logpdf(
                    _w_sample,
                    self.w_priors[idx][0] * (1.0 / brd_samples),
                    self.w_priors[idx][1]
                    * (1.0 / brd_samples)
                    * (1.0 / _wm_sample)
                    * self.layer_sizes[idx],
                )
            elif idx == 0:
                global_pl += gamma_logpdf(
                    _w_sample,
                    self.w_priors[idx][0],
                    self.w_priors[idx][1] * self.layer_sizes[idx],
                )
            else:
                if idx == 1 and self.use_brd:
                    global_pl += gamma_logpdf(
                        _w_sample,
                        self.w_priors[idx][0],
                        self.w_priors[idx][1]
                        * 1.0
                        / brd_samples.T
                        * self.layer_sizes[idx],
                    )
                else:
                    global_pl += gamma_logpdf(
                        _w_sample,
                        self.w_priors[idx][0],
                        self.w_priors[idx][1] * self.layer_sizes[idx],
                    )
            global_en += lognormal_entropy(_w_shape, _w_rate)

            # z
            _z_shape = z_shapes[:, start:end]
            _z_rate = z_rates[:, start:end]

            _z_shape = jax.lax.cond(
                stop_gradients[idx],
                lambda: jax.lax.stop_gradient(_z_shape),
                lambda: _z_shape,
            )
            _z_rate = jax.lax.cond(
                stop_gradients[idx],
                lambda: jax.lax.stop_gradient(_z_rate),
                lambda: _z_rate,
            )
            _z_sample = lognormal_sample(rng, _z_shape, _z_rate)

            # if idx == 0 and self.use_brd:
            #     z_mean = z_mean * _wm_sample.T

            if self.marginalize_alpha:
                alpha = self.alpha
            else:
                alpha = s_sample
            local_pl += jax.lax.cond(
                idx == self.n_layers - 1,
                lambda: gamma_logpdf(_z_sample, self.top_alpha, self.top_alpha),
                lambda: gamma_logpdf(
                    _z_sample,
                    alpha,
                    alpha / (z_mean),
                ),
            )

            local_en += lognormal_entropy(_z_shape, _z_rate)
            z_mean = jnp.einsum("nk,kp->np", _z_sample, _w_sample)

        n_w_sample = _w_sample  # / jnp.sum(_w_sample, axis=1, keepdims=True)
        n_z_sample = _z_sample  # / jnp.sum(_z_sample, axis=1, keepdims=True)
        if self.use_brd:
            n_w_sample = n_w_sample  # * _wm_sample
        mean_bottom = (
            jnp.einsum("nk,kg->ng", n_z_sample, n_w_sample) * cell_budget_sample
        )  # self.exposures[indices]
        mean_bottom = mean_bottom * (batch_indices_onehot.dot(gene_budget_sample))

        # x = jnp.array(batch)
        # lam = mean_bottom  # same shape

        # mask = x > 0
        # x_nz = jnp.where(mask, x, 0.0)

        # # X log λ  -  log(X!)
        # ll_nz = jnp.sum(x_nz * jnp.log(lam + 1e-8))
        # ll_nz -= jnp.sum(jax.scipy.special.gammaln(x_nz + 1.0))

        # # -∑ λ over *all* entries
        # ll_zero = -jnp.sum(lam)

        # ll = ll_nz + ll_zero
        ll = jnp.sum(vmap(poisson.logpmf)(jnp.array(batch), mean_bottom))
        # x = jnp.asarray(batch)
        # x = jnp.maximum(x, 0.0)
        # x = jnp.rint(x).astype(jnp.int32)   # counts should be ints
        # mu = jnp.clip(mean_bottom, 1e-8, 1e8)

        # # phi_sample should be shape (1,G) or (B,G)->(N,G); make sure it broadcasts to mu
        # phi = jnp.clip(s_sample, 1e-3, 1e6)          # IMPORTANT: keep total_count away from 0

        # logits = jnp.log(mu + 1e-8) - jnp.log(phi + 1e-8)

        # nb = tfd.NegativeBinomial(total_count=phi, logits=logits)
        # ll = jnp.sum(nb.log_prob(x))
        ll = jax.lax.cond(
            stop_gradients[0],
            lambda: ll * 0.0,
            lambda: ll,
        )

        # Anneal the entropy
        global_en *= annealing_parameter
        local_en *= annealing_parameter

        return (
            (ll + local_pl + local_en) * (self.X.shape[0] / indices.shape[0])
            + global_pl
            + global_en
        )

    def batch_elbo(
        self,
        rng,
        X,
        indices,
        local_params,
        global_params,
        num_samples,
        annealing_parameter,
        stop_gradients,
        stop_cell_budgets,
        stop_gene_budgets,
    ):
        # Average over a batch of random samples.
        rngs = random.split(rng, num_samples)
        vectorized_elbo = vmap(
            self.elbo, in_axes=(0, None, None, None, None, None, None, None, None)
        )
        return jnp.mean(
            vectorized_elbo(
                rngs,
                X,
                indices,
                local_params,
                global_params,
                annealing_parameter,
                stop_gradients,
                stop_cell_budgets,
                stop_gene_budgets,
            )
        )

    def _optimize(
        self,
        local_update_func,
        global_update_func,
        local_params,
        global_params,
        local_opt_state,
        global_opt_state,
        n_epochs=500,
        batch_size=1,
        annealing_parameter=1.0,
        stop_gradients=None,
        stop_cell_budgets=None,
        stop_gene_budgets=None,
        seed=None,
        min_epochs=100,
        tolerance=1e-5,
        patience=5,
        update_locals=True,
        update_globals=True,
    ):
        if seed is None:
            seed = self.seed

        if stop_gradients is None:
            stop_gradients = np.zeros((self.n_layers))
        if stop_cell_budgets is None:
            stop_cell_budgets = 0.0
        if stop_gene_budgets is None:
            stop_gene_budgets = 0.0

        num_complete_batches, leftover = divmod(self.n_cells, batch_size)
        num_batches = num_complete_batches + bool(leftover)
        self.logger.info(
            f"Each epoch contains {num_batches} batches of size {int(min(batch_size, self.n_cells))}"
        )

        def data_stream():
            rng = np.random.RandomState(0)
            while True:
                perm = rng.permutation(self.n_cells)
                for i in range(num_batches):
                    batch_idx = perm[i * batch_size : (i + 1) * batch_size]
                    yield jnp.array(self.X[batch_idx]), jnp.array(batch_idx)

        batches = data_stream()

        losses = []
        global_grads = 0.0
        annealing_parameter = jnp.array(annealing_parameter)
        stop_gradients = jnp.array(stop_gradients)
        stop_cell_budgets = jnp.array(stop_cell_budgets)
        stop_gene_budgets = jnp.array(stop_gene_budgets)
        rng = random.PRNGKey(seed)
        t = 0
        min_loss = np.inf
        early_stop_counter = 0
        stop_early = False
        pbar = tqdm(range(n_epochs))
        try:
            for epoch in pbar:
                epoch_losses = []
                start_time = time.time()
                for it in range(num_batches):
                    rng, rng_input = random.split(rng)
                    X, indices = next(batches)
                    if update_locals:
                        loss, local_params, local_opt_state = local_update_func(
                            X,
                            indices,
                            t,
                            rng_input,
                            local_params,
                            global_params,
                            local_opt_state,
                            global_opt_state,
                            annealing_parameter,
                            stop_gradients,
                            stop_cell_budgets,
                            stop_gene_budgets,
                        )
                    if update_globals:
                        (
                            loss,
                            global_params,
                            global_opt_state,
                            global_grads,
                        ) = global_update_func(
                            X,
                            indices,
                            t,
                            rng_input,
                            local_params,
                            global_params,
                            local_opt_state,
                            global_opt_state,
                            annealing_parameter,
                            stop_gradients,
                            stop_cell_budgets,
                            stop_gene_budgets,
                        )
                    epoch_losses.append(loss)
                    t += 1
                current_loss = np.mean(epoch_losses)
                losses.append(current_loss)

                if epoch >= min_epochs:
                    if min_loss == np.inf:
                        min_loss = current_loss
                        stop_early = False

                    relative_improvement = (min_loss - current_loss) / np.abs(min_loss)
                    min_loss = min(min_loss, current_loss)

                    if relative_improvement < tolerance:
                        early_stop_counter += 1
                    else:
                        early_stop_counter = 0
                    if early_stop_counter >= patience:
                        stop_early = True

                if stop_early:
                    break

                epoch_time = time.time() - start_time  # noqa: F841
                pbar.set_postfix({"Loss": losses[-1]})
                if epoch >= min_epochs:
                    pbar.set_postfix(
                        {"Loss": losses[-1], "Rel. improvement": relative_improvement}
                    )

            if stop_early:
                self.logger.info(
                    "Relative improvement of "
                    f"{relative_improvement:0.4g} < {tolerance:0.4g} "
                    f"for {patience} step(s) in a row, stopping early."
                )
                if (
                    relative_improvement < 0
                    and np.abs(relative_improvement) > 100 * tolerance
                ):
                    self.logger.info(
                        "It seems like the loss increased significantly. "
                        "Consider using a lower learning rate."
                    )
            else:
                self.logger.info(
                    "Learning stopped before convergence. "
                    "Consider increasing the number of iterations."
                )

        except KeyboardInterrupt:
            self.logger.info("Interrupted learning. Exiting safely...")

        return (
            losses,
            local_params,
            global_params,
            local_opt_state,
            global_opt_state,
            global_grads,
        )

    def fit(
        self,
        nmf_init: bool = True,
        max_cells_init: int = 5000,
        z_init_concentration: float = 1.0,
        n_rounds: int = 1,
        **kwargs: Any,
    ) -> None:
        """Fit scDEF, warm-starting from a previous fit when available.
        Args:
            nmf_init: whether to initialize the model with NMF.
            max_cells_init: maximum number of cells to use for initialization.
            z_init_concentration: concentration parameter of a Gamma distribution to sample the initial z values from. If high coverage, prefer higher values to avoid overfitting early.
            n_rounds: number of rounds to run the optimization.
            **kwargs: additional keyword arguments.

            On the first call, parameters are initialized from priors (or NMF if enabled).
            On subsequent calls, the model is re-initialized from the current posterior
            quantities and the current `factor_lists`, enabling a fit -> filter -> fit
            workflow. During refit, upper-layer sizes are clipped to respect
            ``decay_factor`` before rebuilding the hierarchy.
        """
        if getattr(self, "_has_fit", False):
            # Keep original kept indices before resizing; update_model_size resets factor_lists.
            old_factor_lists = [np.array(f).copy() for f in self.factor_lists]
            layer_sizes = [len(self.factor_lists[i]) for i in range(self.n_layers)]
            # Enforce geometric decay on refit so upper layers cannot exceed the
            # expected size implied by decay_factor.
            for i in range(1, len(layer_sizes)):
                max_allowed = max(
                    1, int(np.floor(layer_sizes[i - 1] / self.decay_factor))
                )
                if layer_sizes[i] > max_allowed:
                    layer_sizes[i] = max_allowed
            self.update_model_size(
                max_n_factors=max(layer_sizes), layer_sizes=layer_sizes
            )
            self.update_model_priors()
            self.logger.info(
                f"Continuing scDEF from previous fit with layer sizes {self.layer_sizes}."
            )
            nmf_init = False
            init_budgets = False
            init_alpha = False
            l0_keep = np.array(old_factor_lists[0], dtype=int)
            init_z = np.array(self.pmeans[f"{self.layer_names[0]}z"])[:, l0_keep]
            init_w = np.array(self.pmeans["L0W"])[l0_keep]
            init_brd = np.array(self.pmeans["brd"])[l0_keep]
            init_ard = np.array(self.pmeans["factor_means"])[l0_keep]
        else:
            init_budgets = True
            init_alpha = True
            init_z = None
            init_w = None
            init_brd = None
            init_ard = None
        self.init_var_params(
            init_budgets=init_budgets,
            init_alpha=init_alpha,
            init_z=init_z,
            init_w=init_w,
            init_brd=init_brd,
            init_ard=init_ard,
            nmf_init=nmf_init,
            max_cells=max_cells_init,
            z_init_concentration=z_init_concentration,
        )
        self.elbos = []
        self.step_sizes = []
        self._learn(n_rounds=n_rounds, **kwargs)
        self._has_fit = True
        self._fit_revision += 1

    def _learn(
        self,
        n_rounds: Optional[int] = 1,
        n_epoch: Optional[int] = 1000,
        lr: Optional[float] = 1e-2,
        local_lr: Optional[float] = 1e-2,
        annealing: Optional[float] = 1.0,
        num_samples: Optional[int] = 100,
        batch_size: Optional[int] = 256,
        layerwise: Optional[bool] = False,
        min_epochs: Optional[int] = 50,
        tolerance: Optional[float] = 1e-5,
        patience: Optional[int] = 50,
        update_locals: Optional[bool] = True,
        update_globals: Optional[bool] = True,
        stop_cell_budgets: Optional[int] = 0,
        stop_gene_budgets: Optional[int] = 0,
        opt_layer: Optional[int] = None,
        filter: Optional[bool] = True,
        annotate: Optional[bool] = True,
        **kwargs,
    ):
        """Fit a variational approximation to the posterior over scDEF parameters.

        Args:
            n_epoch: number of epochs (full passes of the data).
                Can be a list of ints for multi-step learning. ``n_epochs`` is
                also accepted as a backward-compatible alias.
            lr: learning rate.
                Can be a list of floats for multi-step learning.
            annealing: scale factor for the entropy term.
                Can be a list of floats for multi-step learning.
            num_samples: number of Monte Carlo samples to use in the ELBO approximation.
            batch_size: number of data points to use per iteration. If None, uses all.
                Useful for data sets that do not fit in GPU memory.
            layerwise: whether to optimize the model parameters in a step-wise manner:
                first learn only Layer 0 and 1, and then 2, and then 3, and so on. The size of
                the n_epoch or lr schedules will be ignored, only the first value will be used
                and each step will use that n_epoch value.
            min_epochs: minimum number of epochs for early stopping
            tolerance: maximum relative change in loss for early stopping
            patience: number of epochs for which tolerated loss changes must hold for early stopping
            update_locals: whether to optimize the local parameters
            update_globals: whether to optimize the global parameters

        Notes:
            If all optimization steps run with zero epochs (e.g. ``n_epoch=0``),
            no BRD-based auto-filtering is applied at the end of learning.
        """
        if "n_epochs" in kwargs:
            n_epoch = kwargs.pop("n_epochs")

        n_steps = n_rounds
        if layerwise:
            n_steps = self.n_layers

        n_epoch_schedule = [n_epoch] * n_steps
        did_optimize = np.any(np.asarray(n_epoch_schedule, dtype=int) > 0)

        if layerwise:
            lr_schedule = [lr] * n_steps
            local_lr_schedule = [local_lr] * n_steps
        else:
            lr_schedule = [lr * 0.5**step for step in range(n_steps)]
            local_lr_schedule = [local_lr * 0.5**step for step in range(n_steps)]

        annealing_schedule = [annealing] * n_steps

        if batch_size is None:
            batch_size = self.n_cells

        # Set up jitted functions
        stop_gradients = np.ones((self.n_layers))
        layers_to_optimize = np.arange(self.n_layers)

        def objective(
            X,
            indices,
            local_var_params,
            global_var_params,
            key,
            annealing_parameter,
            stop_gradients,
            stop_cell_budgets,
            stop_gene_budgets,
        ):
            return -self.batch_elbo(
                key,
                X,
                indices,
                local_var_params,
                global_var_params,
                num_samples,
                annealing_parameter,
                stop_gradients,
                stop_cell_budgets,
                stop_gene_budgets,
            )  # minimize -ELBO

        def clip_params(
            params, min_mu=-1e10, max_mu=1e4, min_logstd=-1e10, max_logstd=1e1
        ):
            for i in range(len(params))[:-2]:  # skip the last two params (w and s)
                params[i] = (
                    params[i].at[0].set(jnp.clip(params[i][0], min_mu, max_mu))
                )  # mu
                params[i] = (
                    params[i].at[1].set(jnp.clip(params[i][1], min_logstd, max_logstd))
                )  # logstd
            return params

        local_loss_grad = jit(value_and_grad(objective, argnums=2))
        global_loss_grad = jit(value_and_grad(objective, argnums=3))

        for i in range(len(n_epoch_schedule)):
            n_epochs = n_epoch_schedule[i]
            step_size = lr_schedule[i]
            local_step_size = local_lr_schedule[i]
            anneal_param = annealing_schedule[i]
            self.logger.info(f"Initializing optimizer with learning rate {step_size}.")
            if anneal_param != 1:
                self.logger.info(f"Set annealing parameter to {anneal_param}.")

            local_optimizer = optax.adam(local_step_size)
            global_optimizer = optax.adam(step_size)

            if layerwise:
                if opt_layer is not None:
                    if i != opt_layer:
                        continue
                    layers_to_optimize = np.array([i])
                else:
                    layers_to_optimize = np.array([i])
                # layers_to_optimize = np.arange(2 + i)
                # if i > 0:
                #     layers_to_optimize = np.array([2+i-1])
                self.logger.info(f"Optimizing layers {layers_to_optimize}")
            stop_gradients[layers_to_optimize] = 0.0

            def local_update(
                X,
                indices,
                i,
                key,
                local_params,
                global_params,
                local_opt_state,
                global_opt_state,
                annealing_parameter,
                stop_gradients,
                stop_cell_budgets,
                stop_gene_budgets,
            ):
                value, gradient = local_loss_grad(
                    X,
                    indices,
                    local_params,
                    global_params,
                    key,
                    annealing_parameter,
                    stop_gradients,
                    stop_cell_budgets,
                    stop_gene_budgets,
                )
                updates, local_opt_state = local_optimizer.update(
                    gradient, local_opt_state, local_params
                )
                local_params = optax.apply_updates(local_params, updates)
                local_params = clip_params(local_params)
                return value, local_params, local_opt_state

            def global_update(
                X,
                indices,
                i,
                key,
                local_params,
                global_params,
                local_opt_state,
                global_opt_state,
                annealing_parameter,
                stop_gradients,
                stop_cell_budgets,
                stop_gene_budgets,
            ):
                value, gradient = global_loss_grad(
                    X,
                    indices,
                    local_params,
                    global_params,
                    key,
                    annealing_parameter,
                    stop_gradients,
                    stop_cell_budgets,
                    stop_gene_budgets,
                )
                updates, global_opt_state = global_optimizer.update(
                    gradient, global_opt_state, global_params
                )
                global_params = optax.apply_updates(global_params, updates)
                global_params = clip_params(global_params)
                return value, global_params, global_opt_state, gradient

            # local_opt_state = local_opt_init(self.local_params)
            local_opt_state = local_optimizer.init(self.local_params)
            # global_opt_state = global_opt_init(self.global_params)
            global_opt_state = global_optimizer.init(self.global_params)

            (
                losses,
                local_params,
                global_params,
                local_opt_state,
                global_opt_state,
                global_gradients,
            ) = self._optimize(
                local_update,
                global_update,
                self.local_params,
                self.global_params,
                local_opt_state,
                global_opt_state,
                n_epochs=n_epochs,
                batch_size=batch_size,
                annealing_parameter=anneal_param,
                stop_gradients=stop_gradients,
                stop_cell_budgets=stop_cell_budgets,
                stop_gene_budgets=stop_gene_budgets,
                min_epochs=min_epochs,
                tolerance=tolerance,
                patience=patience,
                update_locals=update_locals,
                update_globals=update_globals,
                **kwargs,
            )
            if update_locals:
                self.local_params = local_params
            if update_globals:
                self.global_params = global_params
            self.elbos.append(losses)
            self.step_sizes.append(lr_schedule[i])

        self.set_posterior_means()
        self.set_posterior_variances()
        if self.use_brd and filter and did_optimize:
            self.filter_factors(upper_only=True)
        else:
            if annotate:
                self.make_layercolors(
                    layer_cpal=self.layer_cpal, lightness_mult=self.lightness_mult
                )
                self.annotate_adata()

    def set_posterior_means(self):
        cell_budget_params = self.local_params[0]
        gene_budget_params = self.global_params[0]
        fscale_params = self.global_params[1]
        wm_params = self.global_params[-2]
        z_params = self.local_params[1]

        self.pmeans = {
            "cell_scale": np.array(
                jnp.exp(
                    cell_budget_params[0] + 0.5 * jnp.exp(cell_budget_params[1]) ** 2
                )
            ),
            "gene_scale": np.array(
                jnp.exp(
                    gene_budget_params[0] + 0.5 * jnp.exp(gene_budget_params[1]) ** 2
                )
            ),
            "factor_concentrations": jnp.exp(
                fscale_params[0] + 0.5 * jnp.exp(fscale_params[1]) ** 2
            ),
            "factor_means": jnp.exp(wm_params[0] + 0.5 * jnp.exp(wm_params[1]) ** 2),
        }
        self.pmeans["brd"] = self.pmeans["factor_concentrations"]  # / jnp.maximum(
        #  self.pmeans["factor_means"], 1.0
        #  )
        self.pmeans["wm"] = self.pmeans["factor_means"]
        self.pmeans["ard"] = self.pmeans["wm"]
        for idx in range(self.n_layers):
            start = sum(self.layer_sizes[:idx])
            end = start + self.layer_sizes[idx]
            self.pmeans[f"{self.layer_names[idx]}z"] = jnp.exp(
                z_params[0][:, start:end]
                + 0.5 * jnp.exp(z_params[1][:, start:end]) ** 2
            )
            _w_shape = self.global_params[2 + idx][0]
            _w_rate = self.global_params[2 + idx][1]
            self.pmeans[f"{self.layer_names[idx]}W"] = jnp.exp(
                _w_shape + 0.5 * jnp.exp(_w_rate) ** 2
            )

        self.pmeans["alpha"] = np.exp(
            self.global_params[-1][0] + 0.5 * np.exp(self.global_params[-1][1]) ** 2
        )

    def set_posterior_variances(self):
        cell_budget_params = self.local_params[0]
        gene_budget_params = self.global_params[0]
        fscale_params = self.global_params[1]
        wm_params = self.global_params[-2]
        z_params = self.local_params[1]

        def _lognormal_var(mu, log_sigma):
            sigma2 = np.exp(log_sigma) ** 2
            return (np.exp(sigma2) - 1.0) * np.exp(2.0 * mu + sigma2)

        self.pvars = {
            "cell_scale": np.array(
                _lognormal_var(cell_budget_params[0], cell_budget_params[1])
            ),
            "gene_scale": np.array(
                _lognormal_var(gene_budget_params[0], gene_budget_params[1])
            ),
            "factor_concentrations": np.array(
                _lognormal_var(fscale_params[0], fscale_params[1])
            ),
            "factor_means": np.array(_lognormal_var(wm_params[0], wm_params[1])),
            "alpha": np.array(
                _lognormal_var(self.global_params[-1][0], self.global_params[-1][1])
            ),
        }

        for idx in range(self.n_layers):
            start = sum(self.layer_sizes[:idx])
            end = start + self.layer_sizes[idx]
            self.pvars[f"{self.layer_names[idx]}z"] = np.array(
                _lognormal_var(z_params[0][:, start:end], z_params[1][:, start:end])
            )
            _w_shape = self.global_params[2 + idx][0]
            _w_rate = self.global_params[2 + idx][1]
            self.pvars[f"{self.layer_names[idx]}W"] = np.array(
                _lognormal_var(_w_shape, _w_rate)
            )

    def filter_factors(
        self,
        brd_min: Optional[float] = 1.0,
        ard_min: Optional[float] = 0.001,
        clarity_min: Optional[float] = 0.5,
        min_cells_upper: Optional[float] = 0.001,
        min_cells_lower: Optional[float] = 0.0,
        filter_up: Optional[bool] = True,
        annotate: Optional[bool] = True,
        upper_only: Optional[bool] = False,
    ):
        """Filter our irrelevant factors based on the BRD posterior or the cell attachments.

        Args:
            thres: minimum factor BRD value
            iqr_mult: multiplier of the difference between the third quartile and the median BRD values to set the threshold
            min_cells_upper: minimum number of cells that each factor in upper layers must have attached to it for it to be kept. If between 0 and 1, fraction. Otherwise, absolute value
            min_cells_lower: minimum number of cells that each factor in layer 0 must have attached to it for it to be kept. If between 0 and 1, fraction. Otherwise, absolute value
            filter_up: whether to remove factors in upper layers via inter-layer attachments
            upper_only: whether to only filter factors in upper layers
        """
        if min_cells_upper != 0:
            if min_cells_upper < 1.0:
                min_cells_upper = max(min_cells_upper * self.adata.shape[0], 10)
        if min_cells_lower != 0:
            if min_cells_lower < 1.0:
                min_cells_lower = max(min_cells_lower * self.adata.shape[0], 10)

        new_factor_lists = []
        for i, layer_name in enumerate(self.layer_names):
            if i == 0:
                if upper_only:
                    keep = np.arange(self.layer_sizes[i])
                else:
                    keep = self.get_effective_factors(
                        brd_min=brd_min,
                        ard_min=ard_min,
                        clarity_min=clarity_min,
                        min_cells=min_cells_lower,
                    )
            else:
                assignments = np.argmax(self.pmeans[f"{layer_name}z"], axis=1)
                counts = np.array(
                    [
                        np.count_nonzero(assignments == a)
                        for a in range(self.layer_sizes[i])
                    ]
                )
                keep = np.array(range(self.layer_sizes[i]))[
                    np.where(counts >= min_cells_upper)[0]
                ]
                if filter_up:
                    mat = self.pmeans[f"{layer_name}W"][keep]
                    assignments = []
                    for factor in new_factor_lists[i - 1]:
                        assignments.append(keep[np.argmax(mat[:, factor])])

                    keep = np.unique(
                        list(set(np.unique(assignments)).intersection(keep))
                    )

            if len(keep) == 0:
                self.logger.info(
                    f"No factors in layer {i} satisfy the filtering criterion. Please adjust the filtering parameters."
                    f"Keeping all factors for layer {i} for now."
                )
                keep = np.arange(self.layer_sizes[i])
            new_factor_lists.append(keep)

        self.factor_lists = new_factor_lists
        self.set_factor_names()
        if "factor_obs" in self.adata.uns:
            self.adata.uns["factor_obs"]["technical"] = False

        self.make_layercolors(
            layer_cpal=self.layer_cpal, lightness_mult=self.lightness_mult
        )
        if annotate:
            self.annotate_adata()

    def set_factor_names(self):
        self.factor_names = [
            [
                f"{self.layer_names[idx]}_{str(i)}"
                for i in range(len(self.factor_lists[idx]))
            ]
            for idx in range(self.n_layers)
        ]

    def annotate_adata(self):
        self.adata.obs["cell_scale"] = self.pmeans["cell_scale"]
        if self.n_batches == 1:
            self.adata.var["gene_scale"] = self.pmeans["gene_scale"][0]
            self.logger.info("Updated adata.var: `gene_scale`")
        else:
            for batch_idx in range(self.n_batches):
                name = f"gene_scale_{batch_idx}"
                self.adata.var[name] = self.pmeans["gene_scale"][batch_idx]
                self.logger.info(f"Updated adata.var: `{name}` for batch {batch_idx}.")

        if not getattr(self, "_preserve_factor_names_on_annotate", False):
            self.set_factor_names()
        ranked_genes, ranked_scores = self.get_signatures_dict(
            scores=True, sorted_scores=True
        )

        for idx in range(self.n_layers):
            layer_name = self.layer_names[idx]
            self.adata.obsm[f"X_{layer_name}"] = np.array(
                self.pmeans[f"{layer_name}z"][:, self.factor_lists[idx]]
            )
            assignments = np.argmax(self.adata.obsm[f"X_{layer_name}"], axis=1)
            self.adata.obs[f"{layer_name}"] = [
                self.factor_names[idx][a] for a in assignments
            ]
            # Make sure factor colors in UMAP respect the palette
            factor_colors = [
                matplotlib.colors.to_hex(self.layer_colorpalettes[idx][i])
                for i in range(len(self.factor_lists[idx]))
            ]

            if layer_name == "marker":
                self.adata.obs[f"{layer_name}"] = pd.Categorical(
                    self.adata.obs[f"{layer_name}"]
                )
                sorted_factors = self.adata.obs_vector(f"{layer_name}")
                sorted_colors = []
                for fac in sorted_factors.categories:
                    pos = np.where(np.array(self.factor_names[idx]) == fac)[0][0]
                    col = factor_colors[pos]
                    sorted_colors.append(col)
                    self.adata.uns[f"{layer_name}_colors"] = sorted_colors
            else:
                self.adata.uns[f"{layer_name}_colors"] = factor_colors

            scores_names = [f + "_score" for f in self.factor_names[idx]]
            df = pd.DataFrame(
                self.adata.obsm[f"X_{layer_name}"],
                index=self.adata.obs.index,
                columns=scores_names,
            )
            if scores_names[0] not in self.adata.obs.columns:
                self.adata.obs = pd.concat([self.adata.obs, df], axis=1)
            else:
                self.adata.obs = self.adata.obs.drop(
                    columns=[col for col in self.adata.obs.columns if "score" in col]
                )
                self.adata.obs = pd.concat([self.adata.obs, df], axis=1)

            self.logger.info(
                f"Updated adata.obs with layer {idx}: `{layer_name}` and `{layer_name}_score` for all factors in layer {idx}"
            )
            self.logger.info(f"Updated adata.obsm with layer {idx}: `X_{layer_name}`")

            factor_names = self.factor_names[idx]
            names = np.array(
                [
                    tuple(
                        [ranked_genes[factor_name][i] for factor_name in factor_names]
                    )
                    for i in range(self.n_genes)
                ],
                dtype=[(factor_name, "O") for factor_name in factor_names],
            ).view(np.recarray)

            scores = np.array(
                [
                    tuple(
                        [ranked_scores[factor_name][i] for factor_name in factor_names]
                    )
                    for i in range(self.n_genes)
                ],
                dtype=[(factor_name, "<f4") for factor_name in factor_names],
            ).view(np.recarray)

            pvals = np.array(
                [
                    tuple([1.0 for factor_name in factor_names])
                    for i in range(self.n_genes)
                ],
                dtype=[(factor_name, "<f4") for factor_name in factor_names],
            ).view(np.recarray)

            pvals_adj = np.array(
                [
                    tuple([1.0 for factor_name in factor_names])
                    for i in range(self.n_genes)
                ],
                dtype=[(factor_name, "<f4") for factor_name in factor_names],
            ).view(np.recarray)

            logfoldchanges = np.array(
                [
                    tuple([0.0 for factor_name in factor_names])
                    for i in range(self.n_genes)
                ],
                dtype=[(factor_name, "<f4") for factor_name in factor_names],
            ).view(np.recarray)

            self.adata.uns[f"{layer_name}_signatures"] = {
                "params": {
                    "reference": "rest",
                    "method": "scDEF",
                    "groupby": f"{layer_name}",
                },
                "names": names,
                "scores": scores,
                "pvals": pvals,
                "pvals_adj": pvals_adj,
                "logfoldchanges": logfoldchanges,
            }

            self.logger.info(
                f"Updated adata.uns with layer {idx} signatures: `{layer_name}_signatures`."
            )

        self.normalize_cellscores()

    def normalize_cellscores(self):
        for idx in range(self.n_layers):
            layer_name = self.layer_names[idx]
            self.adata.obsm[f"X_{layer_name}_probs"] = (
                self.adata.obsm[f"X_{layer_name}"]
                / np.sum(self.adata.obsm[f"X_{layer_name}"], axis=1)[:, None]
            )
            scores_names = [f + "_prob" for f in self.factor_names[idx]]
            df = pd.DataFrame(
                self.adata.obsm[f"X_{layer_name}_probs"],
                index=self.adata.obs.index,
                columns=scores_names,
            )
            if scores_names[0] not in self.adata.obs.columns:
                self.adata.obs = pd.concat([self.adata.obs, df], axis=1)
            else:
                self.adata.obs = self.adata.obs.drop(
                    columns=[col for col in self.adata.obs.columns if "prob" in col]
                )
                self.adata.obs = pd.concat([self.adata.obs, df], axis=1)

    def get_annotations(
        self,
        marker_reference: Mapping[str, Sequence[str]],
        gene_rankings: Optional[List[List[str]]] = None,
    ) -> List[List[str]]:
        """Get annotations for factors based on marker gene reference.

        Args:
            marker_reference: dictionary mapping annotation names to gene lists
            gene_rankings: gene rankings for each factor, if None will be computed

        Returns:
            list of annotation lists, one per factor
        """
        if gene_rankings is None:
            gene_rankings = self.get_rankings(layer_idx=0)

        annotations = []
        keys = list(marker_reference.keys())
        for rank in gene_rankings:
            # Get annotation in marker_reference that contains the highest score
            # score is length of intersection over length of union
            scores = np.array(
                [
                    len(set(rank).intersection(set(marker_reference[a])))
                    / len(set(rank).union(set(marker_reference[a])))
                    for a in keys
                ]
            )
            sorted_ann_idx = np.argsort(scores)[::-1]
            sorted_scores = scores[sorted_ann_idx]

            # Get only annotations for which score is not zero
            sorted_ann_idx = sorted_ann_idx[np.where(sorted_scores > 0)[0]]

            ann = np.array(keys)[sorted_ann_idx]
            ann = np.array(
                [f"{a} ({sorted_scores[i]:.4f})" for i, a in enumerate(ann)]
            ).tolist()
            annotations.append(ann)
        return annotations

    def get_rankings(
        self,
        layer_idx: int = 0,
        top_genes: Optional[int] = None,
        genes: bool = True,
        return_scores: bool = False,
        sorted_scores: bool = True,
        drop_factors: Optional[List[str]] = None,
    ) -> Union[List[List[str]], Tuple[List[List[str]], List[List[float]]]]:
        """Get gene or factor rankings for each factor in a layer.

        Args:
            layer_idx: layer index to get rankings for
            top_genes: number of top genes/factors to return
            genes: whether to return gene rankings (True) or factor rankings (False)
            return_scores: whether to return scores along with rankings
            sorted_scores: whether to return scores sorted by ranking
            drop_factors: list of factors to drop from rankings

        Returns:
            list of rankings per factor, or tuple of (rankings, scores) if return_scores is True
        """
        if top_genes is None:
            top_genes = len(self.adata.var_names)

        term_names = np.array(self.adata.var_names)
        factor_names_0 = np.array(self.factor_names[0])
        indices_0 = np.arange(len(factor_names_0))

        # Prepare mask for drop_factors to select factors in dot products for upper layers
        if drop_factors is not None and len(drop_factors) > 0:
            drop_set = set(drop_factors)
            keep_indices_0 = [
                i for i, name in enumerate(factor_names_0) if name not in drop_set
            ]
        else:
            keep_indices_0 = indices_0

        term_scores = self.pmeans[f"{self.layer_names[0]}W"][self.factor_lists[0]]
        n_factors = len(self.factor_lists[layer_idx])

        if layer_idx > 0:
            if genes:
                # Compute factor-to-gene mapping for this upper layer
                term_scores = self.pmeans[f"{self.layer_names[layer_idx]}W"][
                    self.factor_lists[layer_idx]
                ][:, self.factor_lists[layer_idx - 1]]
                for layer in range(layer_idx - 1, 0, -1):
                    lower_mat = self.pmeans[f"{self.layer_names[layer]}W"][
                        self.factor_lists[layer]
                    ][:, self.factor_lists[layer - 1]]
                    term_scores = term_scores.dot(lower_mat)
                # For mapping to genes, drop columns of w0 corresponding to factors in drop_factors
                w0 = self.pmeans[f"{self.layer_names[0]}W"][self.factor_lists[0]]
                if len(keep_indices_0) != len(factor_names_0):
                    w0 = w0[keep_indices_0, :]
                term_scores = (
                    term_scores[:, keep_indices_0]
                    if term_scores.shape[1] == len(factor_names_0)
                    else term_scores
                )
                term_scores = term_scores.dot(w0)
            else:
                n_factors_below = len(self.factor_lists[layer_idx - 1])
                term_names = np.arange(n_factors_below).astype(str)
                term_scores = self.pmeans[f"{self.layer_names[layer_idx]}W"][
                    self.factor_lists[layer_idx]
                ][:, self.factor_lists[layer_idx - 1]]

        top_terms = []
        top_scores = []
        for k in range(n_factors):
            top_terms_idx = (term_scores[k, :]).argsort()[::-1]
            top_terms_idx = top_terms_idx[:top_genes]
            top_terms_list = term_names[top_terms_idx].tolist()
            top_terms.append(top_terms_list)
            if sorted_scores:
                sorted_term_scores_k = term_scores[k, :][top_terms_idx]
                top_scores_list = sorted_term_scores_k.tolist()
            else:
                top_scores_list = term_scores[k, :].tolist()
            top_scores.append(top_scores_list)

        if return_scores:
            return top_terms, top_scores
        return top_terms

    def get_signature_sample(
        self,
        rng: Any,
        factor_idx: int,
        layer_idx: int,
        top_genes: int = 10,
        return_scores: bool = False,
    ) -> Union[List[str], Tuple[List[str], np.ndarray]]:
        """Get a single signature sample from the posterior for a factor.

        Args:
            rng: JAX random number generator key
            factor_idx: index of the factor
            layer_idx: layer index of the factor
            top_genes: number of top genes to return
            return_scores: whether to return scores along with gene names

        Returns:
            list of gene names, or tuple of (gene_names, scores) if return_scores is True
        """
        term_names = np.array(self.adata.var_names)

        term_scores_shape = self.global_params[2 + 0][0][self.factor_lists[0]]
        term_scores_rate = np.exp(self.global_params[2 + 0][1][self.factor_lists[0]])
        rng, sample_rng = random.split(rng)
        term_scores_sample = lognormal_sample(
            sample_rng, term_scores_shape, term_scores_rate
        )

        if layer_idx > 0:
            term_scores_shape = self.global_params[2 + layer_idx][0][
                self.factor_lists[layer_idx]
            ][:, self.factor_lists[layer_idx - 1]]

            term_scores_rate = np.exp(
                self.global_params[2 + layer_idx][1][self.factor_lists[layer_idx]][
                    :, self.factor_lists[layer_idx - 1]
                ]
            )
            rng, sample_rng = random.split(rng)
            term_scores_sample = lognormal_sample(
                sample_rng, term_scores_shape, term_scores_rate
            )

            for layer in range(layer_idx - 1, 0, -1):
                lower_mat_shape = self.global_params[2 + layer][0][
                    self.factor_lists[layer]
                ][:, self.factor_lists[layer - 1]]

                lower_mat_rate = np.exp(
                    self.global_params[2 + layer][1][self.factor_lists[layer]][
                        :, self.factor_lists[layer - 1]
                    ]
                )
                rng, sample_rng = random.split(rng)
                lower_mat_sample = lognormal_sample(
                    sample_rng, lower_mat_shape, lower_mat_rate
                )
                term_scores_sample = term_scores_sample.dot(lower_mat_sample)

            lower_term_scores_shape = self.global_params[2 + 0][0][self.factor_lists[0]]

            lower_term_scores_rate = np.exp(
                self.global_params[2 + 0][1][self.factor_lists[0]]
            )
            rng, sample_rng = random.split(rng)
            lower_term_scores_sample = lognormal_sample(
                sample_rng, lower_term_scores_shape, lower_term_scores_rate
            )

            term_scores_sample = term_scores_sample.dot(lower_term_scores_sample)

        top_terms_idx = (term_scores_sample[factor_idx, :]).argsort()[::-1][:top_genes]
        top_terms = term_names[top_terms_idx].tolist()
        top_scores = term_scores_sample[factor_idx, :].tolist()

        if return_scores:
            return top_terms, top_scores
        return top_terms

    def get_signature_confidence(
        self,
        factor_idx: int,
        layer_idx: int,
        mc_samples: int = 100,
        top_genes: int = 10,
        pairwise: bool = False,
    ) -> float:
        """Get confidence score for a factor signature using Monte Carlo sampling.

        Args:
            factor_idx: index of the factor
            layer_idx: layer index of the factor
            mc_samples: number of Monte Carlo samples to take
            top_genes: number of top genes to consider in each sample
            pairwise: whether to compute pairwise Jaccard similarities

        Returns:
            confidence score as Jaccard similarity
        """
        signatures = []
        base_rng = random.PRNGKey(0)
        for i in range(mc_samples):
            rng = random.fold_in(base_rng, i)
            signature_sample = self.get_signature_sample(
                rng,
                factor_idx=factor_idx,
                layer_idx=layer_idx,
                top_genes=top_genes,
            )
            signatures.append(signature_sample)

        if pairwise:
            jaccs = np.zeros((mc_samples, mc_samples))
            for i in range(mc_samples):
                for j in range(mc_samples):
                    jaccs[i, j] = score_utils.jaccard_similarity(
                        [signatures[i], signatures[j]]
                    )
            return np.mean(jaccs)
        else:
            return score_utils.jaccard_similarity(signatures)

    def get_relevances_dict(self) -> Dict[str, float]:
        """Get dictionary of factor relevance scores.

        Returns:
            dictionary mapping factor names to relevance scores
        """
        relevance_dict = {}
        for layer_idx in range(self.n_layers):
            for factor_idx, factor_name in enumerate(self.factor_names[layer_idx]):
                relevance_dict[factor_name] = self.pmeans["factor_means"][
                    self.factor_lists[layer_idx][factor_idx]
                ]
        return relevance_dict

    def get_sizes_dict(self) -> Dict[str, float]:
        """Get dictionary of factor sizes (number of cells per factor).

        Returns:
            dictionary mapping factor names to cell counts
        """
        sizes_dict = {}
        for layer_idx in range(self.n_layers):
            layer_sizes = self.adata.obs[
                f"{self.layer_names[layer_idx]}"
            ].value_counts()
            for factor_idx, factor_name in enumerate(self.factor_names[layer_idx]):
                if factor_name in layer_sizes.keys():
                    sizes_dict[factor_name] = layer_sizes[factor_name]
                else:
                    sizes_dict[factor_name] = 0
        return sizes_dict

    def get_signatures_dict(
        self,
        top_genes: Optional[int] = None,
        scores: bool = False,
        sorted_scores: bool = False,
        layer_normalize: bool = False,
        drop_factors: Optional[List[str]] = None,
    ) -> Union[
        Dict[str, List[str]], Tuple[Dict[str, List[str]], Dict[str, np.ndarray]]
    ]:
        """Get dictionary of gene signatures for all factors across all layers.

        Args:
            top_genes: number of top genes per signature
            scores: whether to return scores along with signatures
            sorted_scores: whether to return scores sorted by ranking
            layer_normalize: whether to normalize scores within each layer
            drop_factors: list of factors to exclude

        Returns:
            dictionary mapping factor names to gene lists, or tuple of (signatures, scores) if scores is True
        """
        signatures_dict = {}
        scores_dict = {}
        for layer_idx in range(self.n_layers):
            layer_signatures, layer_scores = self.get_rankings(
                layer_idx=layer_idx,
                top_genes=top_genes,
                return_scores=True,
                sorted_scores=sorted_scores,
                drop_factors=drop_factors,
            )
            for factor_idx, factor_name in enumerate(self.factor_names[layer_idx]):
                val = np.array(layer_scores[factor_idx])
                if layer_normalize:
                    val = val - np.min(val)
                    val = val / np.max(val)
                signatures_dict[factor_name] = layer_signatures[factor_idx]
                scores_dict[factor_name] = val

        if scores:
            return signatures_dict, scores_dict
        return signatures_dict

    def get_summary(self, top_genes: int = 10, reindex: bool = True) -> str:
        """Get a text summary of the model factors and their top genes.

        Args:
            top_genes: number of top genes to show per factor
            reindex: whether to reindex factors

        Returns:
            string summary of the model
        """
        tokeep = self.factor_lists[0]
        n_factors_eff = len(tokeep)
        genes, scores = self.get_rankings(return_scores=True)

        summary = f"Found {n_factors_eff} factors grouped in the following way:\n"

        # Group the factors
        assignments = []
        for i, factor in enumerate(tokeep):
            assignments.append(
                np.argmax(self.pmeans[f"{self.layer_names[1]}W"][:, factor])
            )
        assignments = np.array(assignments)
        factor_order = np.argsort(np.array(assignments))  # noqa: F841

        for group in np.unique(assignments):
            factors = tokeep[np.where(assignments == group)[0]]
            if len(factors) > 0:
                summary += f"Group {group}:\n"
            for i, factor in enumerate(factors):
                summary += f"    Factor {i}: "
                summary += ", ".join(
                    [
                        f"{genes[factor][j]} ({scores[factor][j]:.3f})"
                        for j in range(top_genes)
                    ]
                )
                summary += "\n"
            summary += "\n"

        self.logger.info(summary)

        return summary

    def get_enrichments(
        self,
        libs: List[str] = ["KEGG_2019_Human"],
        gene_rankings: Optional[List[List[str]]] = None,
    ) -> List[Any]:
        """Get gene set enrichments for factor signatures using gseapy.

        Args:
            libs: list of gene set library names to use
            gene_rankings: gene rankings for each factor, if None will be computed

        Returns:
            list of enrichment results, one per factor
        """
        import gseapy as gp

        if gene_rankings is None:
            gene_rankings = self.get_rankings(layer_idx=0)

        enrichments = []
        for rank in tqdm(gene_rankings):
            enr = gp.enrichr(
                gene_list=rank,
                gene_sets=libs,
                organism="Human",
                outdir="test/enrichr",
                cutoff=0.05,
            )
            enrichments.append(enr)
        return enrichments

    def get_layer_factor_orders(self) -> List[np.ndarray]:
        """Get the ordering of factors in each layer for plotting.

        Returns:
            list of arrays, one per layer, containing factor indices in plotting order
        """
        layer_factor_orders = []
        for layer_idx in np.arange(0, self.n_layers)[::-1]:  # Go top down
            factors = self.factor_lists[layer_idx]
            n_factors = len(factors)
            if layer_idx < self.n_layers - 1:
                # Assign factors to upper factors to set the plotting order
                mat = self.pmeans[f"{self.layer_names[layer_idx+1]}W"][
                    self.factor_lists[layer_idx + 1]
                ][:, self.factor_lists[layer_idx]]
                normalized_factor_weights = mat / np.sum(mat, axis=1).reshape(-1, 1)
                assignments = []
                for factor_idx in range(n_factors):
                    assignments.append(
                        np.argmax(normalized_factor_weights[:, factor_idx])
                    )
                assignments = np.array(assignments)

                factor_order = []
                for upper_factor_idx in layer_factor_orders[-1]:
                    factor_order.append(np.where(assignments == upper_factor_idx)[0])
                factor_order = np.concatenate(factor_order).astype(int)
                layer_factor_orders.append(factor_order)
            else:
                layer_factor_orders.append(np.arange(n_factors))
        layer_factor_orders = layer_factor_orders[::-1]
        return layer_factor_orders

    def attach_factors_to_obs(self, obs_key: str) -> List[List[str]]:
        """Attach factors to observation categories.

        Args:
            obs_key: key in model.adata.obs to use for attachment

        Returns:
            list of attachment lists, one per layer
        """
        attachments = []
        for layer, layer_name in enumerate(self.layer_names):
            layer_attachments = []
            for factor_idx in range(len(self.factor_lists[layer])):
                factor_name = f"{self.factor_names[layer][int(factor_idx)]}"
                # cells attached to this factor
                cells = np.where(self.adata.obs[f"{layer_name}"] == factor_name)[0]
                if len(cells) > 0:
                    # cells in this factor that belong to each obs
                    prevs = [
                        np.count_nonzero(self.adata.obs[obs_key][cells] == b)
                        / len(np.where(self.adata.obs[obs_key] == b)[0])
                        for b in self.adata.obs[obs_key].cat.categories
                    ]
                    obs_idx = np.argmax(prevs)  # obs attachment
                    layer_attachments.append(
                        self.adata.obs[obs_key].cat.categories[obs_idx]
                    )
            attachments.append(layer_attachments)
        return attachments

    def compute_weight(self, upper_factor_name: str, lower_factor_name: str) -> float:
        """Compute the weight between two factors across any number of layers.

        Args:
            upper_factor_name: name of the upper factor
            lower_factor_name: name of the lower factor

        Returns:
            weight value between the two factors
        """
        upper_factor_idx = -1
        upper_factor_layer_idx = -1
        for layer_idx in range(self.n_layers):
            layer_factor_names = np.array(self.factor_names[layer_idx])
            if upper_factor_name in layer_factor_names:
                upper_factor_idx = np.where(upper_factor_name == layer_factor_names)[0][
                    0
                ]
                upper_factor_layer_idx = layer_idx
                break

        assert upper_factor_idx != -1

        lower_factor_idx = -1
        lower_factor_layer_idx = -1
        for layer_idx in range(self.n_layers):
            layer_factor_names = np.array(self.factor_names[layer_idx])
            if lower_factor_name in layer_factor_names:
                lower_factor_idx = np.where(lower_factor_name == layer_factor_names)[0][
                    0
                ]
                lower_factor_layer_idx = layer_idx
                break

        assert lower_factor_idx != -1

        upper_layer_name = self.layer_names[upper_factor_layer_idx]
        mat = self.pmeans[f"{upper_layer_name}W"][
            self.factor_lists[upper_factor_layer_idx]
        ][:, self.factor_lists[upper_factor_layer_idx - 1]]
        for layer_idx in range(upper_factor_layer_idx - 1, lower_factor_layer_idx, -1):
            layer_name = self.layer_names[layer_idx]
            lower_mat = self.pmeans[f"{layer_name}W"][self.factor_lists[layer_idx]][
                :, self.factor_lists[layer_idx - 1]
            ]
            mat = mat.dot(lower_mat)

        return mat[upper_factor_idx][lower_factor_idx]

    def compute_factor_obs_association_score(
        self, layer_idx: int, factor_name: str, obs_key: str, obs_val: str
    ) -> float:
        """Compute association score between a factor and an observation category.

        Args:
            layer_idx: layer index of the factor
            factor_name: name of the factor
            obs_key: key in model.adata.obs
            obs_val: value in obs_key to compute association with

        Returns:
            association score value
        """
        layer_name = self.layer_names[layer_idx]

        # Cells attached to factor
        adata_cells_in_factor = self.adata[
            np.where(self.adata.obs[f"{layer_name}"] == factor_name)[0]
        ]

        # Cells from obs_val
        adata_cells_from_obs = self.adata[
            np.where(self.adata.obs[obs_key] == obs_val)[0]
        ]

        cells_from_obs = float(adata_cells_from_obs.shape[0])  # noqa: F841

        # Number of cells from obs_val that are not in factor
        cells_not_in_factor_from_obs = float(
            np.count_nonzero(adata_cells_from_obs.obs[f"{layer_name}"] != factor_name)
        )

        # Number of cells in factor that are obs_val
        cells_in_factor_from_obs = float(
            np.count_nonzero(adata_cells_in_factor.obs[obs_key] == obs_val)
        )

        # Number of cells in factor that are not obs_val
        cells_in_factor_not_from_obs = float(
            np.count_nonzero(adata_cells_in_factor.obs[obs_key] != obs_val)
        )

        return score_utils.compute_fscore(
            cells_in_factor_from_obs,
            cells_in_factor_not_from_obs,
            cells_not_in_factor_from_obs,
        )

    def compute_factor_obs_assignment_fracs(
        self, layer_idx: int, factor_name: str, obs_key: str, obs_val: str
    ) -> float:
        """Compute assignment fraction between a factor and an observation category.

        Args:
            layer_idx: layer index of the factor
            factor_name: name of the factor
            obs_key: key in model.adata.obs
            obs_val: value in obs_key to compute fraction with

        Returns:
            assignment fraction value
        """
        layer_name = self.layer_names[layer_idx]

        # Cells attached to factor
        adata_cells_in_factor = self.adata[
            np.where(self.adata.obs[f"{layer_name}"] == factor_name)[0]
        ]

        # Cells in factor
        cells_in_factor = float(adata_cells_in_factor.shape[0])

        # Cells from factor in obs
        cells_in_factor_from_obs = float(
            np.count_nonzero(adata_cells_in_factor.obs[obs_key] == obs_val)
        )

        score = 0.0
        if cells_in_factor != 0:
            score = cells_in_factor_from_obs / cells_in_factor

        return score

    def compute_factor_obs_weight_score(
        self, layer_idx: int, factor_name: str, obs_key: str, obs_val: str
    ) -> float:
        """Compute weight score between a factor and an observation category.

        Args:
            layer_idx: layer index of the factor
            factor_name: name of the factor
            obs_key: key in model.adata.obs
            obs_val: value in obs_key to compute weight with

        Returns:
            weight score value
        """
        layer_name = self.layer_names[layer_idx]  # noqa: F841

        # Cells from obs_val
        adata_cells_from_obs = self.adata[
            np.where(self.adata.obs[obs_key] == obs_val)[0]
        ]
        adata_cells_not_from_obs = self.adata[
            np.where(self.adata.obs[obs_key] != obs_val)[0]
        ]

        # Weight of cells from obs in factor
        avg_in = np.mean(adata_cells_from_obs.obs[f"{factor_name}_score"])

        # Weight of cells not from obs in factor
        avg_out = np.mean(adata_cells_not_from_obs.obs[f"{factor_name}_score"])

        score = avg_in / np.sum(avg_in + avg_out)

        return score
