"""
Dirichlet-multinomial models for statistical analysis of compositional changes in single-cell data.

For further reference, see:
Büttner, Ostner et al.: scCODA: A Bayesian model for compositional single-cell data analysis

:authors: Johannes Ostner
"""
import numpy as np
import time
import warnings

import tensorflow as tf
import tensorflow_probability as tfp

from sccoda.util import result_classes as res
from typing import Optional, Tuple, Collection, Union, List

tfd = tfp.distributions
tfb = tfp.bijectors


class CompositionalModel:
    """
    Dynamical framework for formulation and inference of Bayesian models for compositional data analysis.
    This framework is used to implement scCODA's model as a subclass of this class.
    Tensorflow probability then allows to run a multitude of inference algorithms on these models
    without the need to specify them every time.

    A `CompositionalModel` consists of the following parameters:

    - `covariate_matrix`: Numpy array that specifies the independent variables (X). Generated equivalently to the covariate matrix of a linear regression.

    - `data_matrix`: Dependent variable (Y). Includes the raw cell type counts for every sample.

    - `cell_types`, covariate_names: Names of cell types and covariates

    - `formula`: String that represents which covariates to include and how to transform them. Used analogously to the formula in R's lm function

    A specific compositional model is then implemented as a child class, with the following additional parameters
    specified in the constructor:

    - `target_log_prob_fn`: Log-probability function of the model. For more specific information, please refer to (tensorflow probability's API)[https://www.tensorflow.org/probability/api_docs/python/tfp]

    - `param_names`: Names of prior and intermediate parameters that are included in the model output. The order has to be the same as in the states_burnin output of `self.get_y_hat`

    - `init_params`: Initial values for the inference method

    Methods implemented by this class:

    - `sampling`: General MCMC sampling that uses a transition kernel

    - `get_chains_after_burnin`: Application of burn-in to MCMC sampling results

    - MCMC sampling methods (`sample_hmc`, `sample_hmc_da`, `sample_nuts`)

    Methods implemented by a child class:

    - `get_y_hat`: Calculation of intermediate parameters for all MCMC chain states and posterior mode of the cell count matrix

    """

    def __init__(
            self,
            covariate_matrix: np.ndarray,
            data_matrix: np.ndarray,
            cell_types: List[str],
            covariate_names: List[str],
            formula: str,
            *args,
            **kwargs
    ):
        """
        Generalized Constructor of Bayesian compositional model class.

        Parameters
        ----------
        covariate_matrix
            covariate matrix, size NxD
        data_matrix
            cell count matrix, size NxK
        cell_types
            Cell type names
        covariate_names
            Covariate names
        """

        dtype = tf.float64
        self.x = tf.convert_to_tensor(covariate_matrix, dtype)

        # Add pseudocount if zeroes are present.
        if np.count_nonzero(data_matrix) != np.size(data_matrix):
            print("Zero counts encountered in data! Added a pseudocount of 0.5.")
            data_matrix += 0.5
        self.y = tf.convert_to_tensor(data_matrix, dtype)

        sample_counts = np.sum(data_matrix, axis=1)
        self.n_total = tf.cast(sample_counts, dtype)
        self.cell_types = cell_types
        self.covariate_names = covariate_names
        self.formula = formula

        # Get dimensions of data
        self.N, self.D = self.x.shape
        self.K = self.y.shape[1]

        # Check input data
        if self.N != self.y.shape[0]:
            raise ValueError("Wrong input dimensions X[{},:] != y[{},:]".format(self.x.shape[0], self.y.shape[0]))
        if self.N != len(self.n_total):
            raise ValueError("Wrong input dimensions X[{},:] != n_total[{}]".format(self.x.shape[0], len(self.n_total)))

    def sampling(
            self,
            num_results: int,
            num_burnin: int,
            kernel,
            init_state: dict,
            trace_fn,
    ) -> Tuple[List[any], List[any], float]:
        """
        MCMC sampling process (tensorflow 2)

        Parameters
        ----------
        num_results
            MCMC chain length (default 20000)
        num_burnin
            Number of burnin iterations (default 5000)
        kernel
            tensorflow MCMC kernel object
        init_state
            Starting parameters
        trace_fn
            tracing function

        Returns
        -------
        MCMC chain states and results

        states
            States of MCMC chain
        kernel_results
            sampling meta-information
        duration
            Duration of MCMC sampling process
        """

        # HMC sampling function
        @tf.function
        def sample_mcmc(num_results_, num_burnin_, kernel_, current_state_, trace_fn):

            return tfp.mcmc.sample_chain(
                num_results=num_results_,
                num_burnin_steps=num_burnin_,
                kernel=kernel_,
                current_state=current_state_,
                trace_fn=trace_fn
            )

        # The actual sampling process
        start = time.time()
        states, kernel_results = sample_mcmc(num_results, num_burnin, kernel, init_state, trace_fn)
        duration = time.time() - start
        print("MCMC sampling finished. ({:.3f} sec)".format(duration))

        return states, kernel_results, duration

    def get_chains_after_burnin(
            self,
            samples: List[any],
            kernel_results: List[any],
            num_burnin: int,
            is_nuts: bool = False
    ) -> Tuple[List[any], dict, float]:
        """
        Application of burn-in after MCMC sampling.
        Cuts the first `num_burnin` samples from all inferred variables and diagnostic statistics.

        Parameters
        ----------
        samples
            all kernel states
        kernel_results
            Kernel meta-information. The tracked statistics depend on the sampling method.
        num_burnin
            number of burn-in iterations
        is_nuts
            Specifies whether NUTS sampling was used

        Returns
        -------
        MCMC chain without burn-in, sampling statistics, acceptance rate

        states_burnin
            Kernel states without burn-in samples
        stats
            sampling statistics
        p_accept
            acceptance rate of MCMC process
        """

        # Samples after burn-in
        states_burnin = []
        stats = {}

        # Apply burn-in to MCMC results
        for s in samples:
            states_burnin.append(s[num_burnin:].numpy())

        # Apply burn-in to sampling statistics
        for k, v in kernel_results.items():
            stats[k] = v[num_burnin:].numpy()

        # Calculation of acceptance rate (different for NUTS sampling)
        if is_nuts:
            p_accept = np.mean(np.exp(kernel_results["log_accept_ratio"].numpy()))
        else:
            acceptances = kernel_results["is_accepted"].numpy()

            # Calculate acceptance rate
            p_accept = sum(acceptances) / acceptances.shape[0]
        print('Acceptance rate: %0.1f%%' % (100 * p_accept))

        return states_burnin, stats, p_accept

    def sample_hmc(
            self,
            num_results: int = int(20e3),
            num_burnin: int = int(5e3),
            num_adapt_steps: Optional[int] = None,
            num_leapfrog_steps: Optional[int] = 10,
            step_size: float = 0.01
    ) -> res.CAResult:

        """
        Hamiltonian Monte Carlo (HMC) sampling in tensorflow 2.

        Tracked diagnostic statistics:

        - `target_log_prob`: Value of the model's log-probability

        - `diverging`: Marks samples as diverging (NOTE: Handle with care, the spike-and-slab prior of scCODA usually leads to many samples being flagged as diverging)

        - `is_accepted`: Whether the proposed sample was accepted in the algorithm's acceptance step

        - `step_size`: The step size used by the algorithm in each step

        Parameters
        ----------
        num_results
            MCMC chain length (default 20000)
        num_burnin
            Number of burnin iterations (default 5000)
        num_adapt_steps
            Length of step size adaptation procedure
        num_leapfrog_steps
            HMC leapfrog steps (default 10)
        step_size
            Initial step size (default 0.01)

        Returns
        -------
        results object

        result
            Compositional analysis result
        """

        # bijectors (not in use atm, therefore identity)
        constraining_bijectors = [tfb.Identity() for x in range(len(self.init_params))]

        # HMC transition kernel
        hmc_kernel = tfp.mcmc.HamiltonianMonteCarlo(
            target_log_prob_fn=self.target_log_prob_fn,
            step_size=step_size,
            num_leapfrog_steps=num_leapfrog_steps)
        hmc_kernel = tfp.mcmc.TransformedTransitionKernel(
            inner_kernel=hmc_kernel, bijector=constraining_bijectors)

        # Set default value for adaptation steps if none given
        if num_adapt_steps is None:
            num_adapt_steps = int(0.8 * num_burnin)

        # Add step size adaptation (Andrieu, Thomas - 2008)
        hmc_kernel = tfp.mcmc.SimpleStepSizeAdaptation(
            inner_kernel=hmc_kernel, num_adaptation_steps=num_adapt_steps, target_accept_prob=0.8)

        # diagnostics tracing function
        def trace_fn(_, pkr):
            return {
                'target_log_prob': pkr.inner_results.inner_results.accepted_results.target_log_prob,
                'diverging': (pkr.inner_results.inner_results.log_accept_ratio < -1000.),
                'is_accepted': pkr.inner_results.inner_results.is_accepted,
                'step_size': pkr.inner_results.inner_results.accepted_results.step_size,
            }

        # The actual HMC sampling process
        states, kernel_results, duration = self.sampling(num_results, num_burnin,
                                                         hmc_kernel, self.init_params, trace_fn)

        # apply burn-in
        states_burnin, sample_stats, acc_rate = self.get_chains_after_burnin(states, kernel_results, num_burnin,
                                                                             is_nuts=False)

        # Calculate posterior predictive
        y_hat = self.get_y_hat(states_burnin, num_results, num_burnin)

        params = dict(zip(self.param_names, states_burnin))

        # Result object generation setup
        # Get names of cell types that are not the reference
        if self.reference_cell_type is not None:
            cell_types_nb = self.cell_types[:self.reference_cell_type] + self.cell_types[self.reference_cell_type+1:]
        else:
            cell_types_nb = self.cell_types

        # Result object generation process. Uses arviz's data structure.
        posterior = {var_name: [var] for var_name, var in params.items() if
                     "prediction" not in var_name}

        if "prediction" in self.param_names:
            posterior_predictive = {"prediction": [params["prediction"]]}
        else:
            posterior_predictive = {}

        observed_data = {"y": self.y}
        dims = {"alpha": ["cell_type"],
                "mu_b": ["1"],
                "sigma_b": ["1"],
                "b_offset": ["covariate", "cell_type_nb"],
                "ind_raw": ["covariate", "cell_type_nb"],
                "ind": ["covariate", "cell_type_nb"],
                "b_raw": ["covariate", "cell_type_nb"],
                "beta": ["covariate", "cell_type"],
                "concentration": ["sample", "cell_type"],
                "prediction": ["sample", "cell_type"]
                }
        coords = {"cell_type": self.cell_types,
                  "cell_type_nb": cell_types_nb,
                  "covariate": self.covariate_names,
                  "sample": range(self.y.shape[0])
                  }

        sampling_stats = {"chain_length": num_results, "num_burnin": num_burnin,
                          "acc_rate": acc_rate, "duration": duration, "y_hat": y_hat}

        model_specs = {"reference": self.reference_cell_type, "formula": self.formula}

        return res.CAResultConverter(posterior=posterior,
                                     posterior_predictive=posterior_predictive,
                                     observed_data=observed_data,
                                     dims=dims,
                                     sample_stats=sample_stats,
                                     coords=coords).to_result_data(sampling_stats=sampling_stats,
                                                                   model_specs=model_specs)

    def sample_hmc_da(
            self,
            num_results: int = int(20e3),
            num_burnin: int = int(5e3),
            num_adapt_steps: Optional[int] = None,
            num_leapfrog_steps: Optional[int] = 10,
            step_size: float = 0.01
    ) -> res.CAResult:
        """
        HMC sampling with dual-averaging step size adaptation (Nesterov, 2009)

        Tracked diagnostic statistics:

        - `target_log_prob`: Value of the model's log-probability

        - `diverging`: Marks samples as diverging (NOTE: Handle with care, the spike-and-slab prior of scCODA usually leads to many samples being flagged as diverging)

        - `log_acc_ratio`: log-acceptance ratio

        - `is_accepted`: Whether the proposed sample was accepted in the algorithm's acceptance step

        - `step_size`: The step size used by the algorithm in each step

        Parameters
        ----------
        num_results
            MCMC chain length (default 20000)
        num_burnin
            Number of burnin iterations (default 5000)
        num_adapt_steps
            Length of step size adaptation procedure
        num_leapfrog_steps
            HMC leapfrog steps (default 10)
        step_size
            Initial step size (default 0.01)

        Returns
        -------
        result object

        result
            Compositional analysis result
        """

        warnings.warn(
            "This feature is untested and might yield different results than expected. Please use sample_hmc().",
            category=UserWarning
        )

        # bijectors (not in use atm, therefore identity)
        constraining_bijectors = [tfb.Identity() for x in range(len(self.init_params))]

        # HMC transition kernel
        hmc_kernel = tfp.mcmc.HamiltonianMonteCarlo(
            target_log_prob_fn=self.target_log_prob_fn,
            step_size=step_size,
            num_leapfrog_steps=num_leapfrog_steps)
        hmc_kernel = tfp.mcmc.TransformedTransitionKernel(
            inner_kernel=hmc_kernel, bijector=constraining_bijectors)

        # Set default value for adaptation steps if none given
        if num_adapt_steps is None:
            num_adapt_steps = int(0.8 * num_burnin)

        # Add step size adaptation
        hmc_kernel = tfp.mcmc.DualAveragingStepSizeAdaptation(
            inner_kernel=hmc_kernel, num_adaptation_steps=num_adapt_steps, target_accept_prob=0.8, decay_rate=0.5)

        # tracing function
        def trace_fn(_, pkr):
            return {
                'target_log_prob': pkr.inner_results.inner_results.accepted_results.target_log_prob,
                'diverging': (pkr.inner_results.inner_results.log_accept_ratio < -1000.),
                "log_acc_ratio": pkr.inner_results.inner_results.log_accept_ratio,
                'is_accepted': pkr.inner_results.inner_results.is_accepted,
                'step_size': tf.exp(pkr.log_averaging_step[0]),
            }

        # HMC sampling
        states, kernel_results, duration = self.sampling(num_results, num_burnin, hmc_kernel, self.init_params, trace_fn)
        states_burnin, sample_stats, acc_rate = self.get_chains_after_burnin(states, kernel_results, num_burnin,
                                                                             is_nuts=False)

        y_hat = self.get_y_hat(states_burnin, num_results, num_burnin)

        params = dict(zip(self.param_names, states_burnin))

        # Specification of cell types that were not used as the reference
        if self.reference_cell_type is not None:
            cell_types_nb = self.cell_types[:self.reference_cell_type] + self.cell_types[self.reference_cell_type+1:]
        else:
            cell_types_nb = self.cell_types

        # Result object generation process. Uses arviz's data structure.
        posterior = {var_name: [var] for var_name, var in params.items() if
                     "prediction" not in var_name}

        if "prediction" in self.param_names:
            posterior_predictive = {"prediction": [params["prediction"]]}
        else:
            posterior_predictive = {}

        observed_data = {"y": self.y}
        dims = {"alpha": ["cell_type"],
                "mu_b": ["1"],
                "sigma_b": ["1"],
                "b_offset": ["covariate", "cell_type_nb"],
                "ind_raw": ["covariate", "cell_type_nb"],
                "ind": ["covariate", "cell_type_nb"],
                "b_raw": ["covariate", "cell_type_nb"],
                "beta": ["covariate", "cell_type"],
                "concentration": ["sample", "cell_type"],
                "prediction": ["sample", "cell_type"]
                }
        coords = {"cell_type": self.cell_types,
                  "cell_type_nb": cell_types_nb,
                  "covariate": self.covariate_names,
                  "sample": range(self.y.shape[0])
                  }

        # build dictionary with sampling statistics
        sampling_stats = {"chain_length": num_results, "num_burnin": num_burnin,
                          "acc_rate": acc_rate, "duration": duration, "y_hat": y_hat}

        model_specs = {"reference": self.reference_cell_type, "formula": self.formula}

        return res.CAResultConverter(posterior=posterior,
                                     posterior_predictive=posterior_predictive,
                                     observed_data=observed_data,
                                     dims=dims,
                                     sample_stats=sample_stats,
                                     coords=coords).to_result_data(sampling_stats=sampling_stats,
                                                                   model_specs=model_specs)

    def sample_nuts(
            self,
            num_results: int = int(10e3),
            num_burnin: int = int(5e3),
            num_adapt_steps: Optional[int] = None,
            max_tree_depth: int = 10,
            step_size: float = 0.01
    ) -> res.CAResult:
        """
        HMC with No-U-turn (NUTS) sampling.
        This method is untested and might yield different results than expected.

        Tracked diagnostic statistics:

        - `target_log_prob`: Value of the model's log-probability

        - `leapfros_taken`: Number of leapfrog steps taken by hte integrator

        - `diverging`: Marks samples as diverging (NOTE: Handle with care, the spike-and-slab prior of scCODA usually leads to many samples being flagged as diverging)

        - `energy`: HMC "Energy" value for each step

        - `log_accept_ratio`: log-acceptance ratio

        - `step_size`: The step size used by the algorithm in each step

        - `reached_max_depth`: Whether the NUTS algorithm reached the maximum sampling depth in each step

        - `is_accepted`: Whether the proposed sample was accepted in the algorithm's acceptance step

        Parameters
        ----------
        num_results
            MCMC chain length (default 10000)
        num_burnin
            Number of burnin iterations (default 5000)
        num_adapt_steps
            Length of step size adaptation procedure
        max_tree_depth
            Maximum tree depth (default 10)
        step_size
            Initial step size (default 0.01)

        Returns
        -------
        result object

        result
            Compositional analysis result
        """

        warnings.warn(
            "This feature is untested and might yield different results than expected. Please use sample_hmc().",
            category=UserWarning
        )

        # bijectors (not in use atm, therefore identity)
        constraining_bijectors = [tfb.Identity() for x in range(len(self.init_params))]

        # NUTS transition kernel
        nuts_kernel = tfp.mcmc.NoUTurnSampler(
            target_log_prob_fn=self.target_log_prob_fn,
            step_size=step_size,
            max_tree_depth=max_tree_depth)
        nuts_kernel = tfp.mcmc.TransformedTransitionKernel(
            inner_kernel=nuts_kernel,
            bijector=constraining_bijectors
        )

        # Set default value for adaptation steps
        if num_adapt_steps is None:
            num_adapt_steps = int(0.8 * num_burnin)

        # Step size adaptation (Nesterov, 2009)
        nuts_kernel = tfp.mcmc.DualAveragingStepSizeAdaptation(
            inner_kernel=nuts_kernel, num_adaptation_steps=num_adapt_steps, target_accept_prob=0.8,
            step_size_setter_fn=lambda pkr, new_step_size: pkr._replace(
                inner_results=pkr.inner_results._replace(step_size=new_step_size)
            ),
            step_size_getter_fn=lambda pkr: pkr.inner_results.step_size,
            log_accept_prob_getter_fn=lambda pkr: pkr.inner_results.log_accept_ratio,
        )

        # trace function
        def trace_fn(_, pkr):
            return {
                "target_log_prob": pkr.inner_results.inner_results.target_log_prob,
                "leapfrogs_taken": pkr.inner_results.inner_results.leapfrogs_taken,
                "diverging": pkr.inner_results.inner_results.has_divergence,
                "energy": pkr.inner_results.inner_results.energy,
                "log_accept_ratio": pkr.inner_results.inner_results.log_accept_ratio,
                "step_size": pkr.inner_results.inner_results.step_size[0],
                "reach_max_depth": pkr.inner_results.inner_results.reach_max_depth,
                "is_accepted": pkr.inner_results.inner_results.is_accepted,
            }

        # HMC sampling
        states, kernel_results, duration = self.sampling(num_results, num_burnin, nuts_kernel, self.init_params, trace_fn)
        states_burnin, sample_stats, acc_rate = self.get_chains_after_burnin(states, kernel_results, num_burnin,
                                                                             is_nuts=True)

        y_hat = self.get_y_hat(states_burnin, num_results, num_burnin)

        params = dict(zip(self.param_names, states_burnin))

        # Result object generation process. Uses arviz's data structure.
        # Get names of cell types that are not the reference
        if self.reference_cell_type is not None:
            cell_types_nb = self.cell_types[:self.reference_cell_type] + self.cell_types[self.reference_cell_type + 1:]
        else:
            cell_types_nb = self.cell_types

        posterior = {var_name: [var] for var_name, var in params.items() if
                     "prediction" not in var_name}

        if "prediction" in self.param_names:
            posterior_predictive = {"prediction": [params["prediction"]]}
        else:
            posterior_predictive = {}

        observed_data = {"y": self.y}
        dims = {"alpha": ["cell_type"],
                "mu_b": ["1"],
                "sigma_b": ["1"],
                "b_offset": ["covariate", "cell_type_nb"],
                "ind_raw": ["covariate", "cell_type_nb"],
                "ind": ["covariate", "cell_type_nb"],
                "b_raw": ["covariate", "cell_type_nb"],
                "beta": ["covariate", "cell_type"],
                "concentration": ["sample", "cell_type"],
                "prediction": ["sample", "cell_type"]
                }
        coords = {"cell_type": self.cell_types,
                  "cell_type_nb": cell_types_nb,
                  "covariate": self.covariate_names,
                  "sample": range(self.y.shape[0])
                  }

        sampling_stats = {"chain_length": num_results, "num_burnin": num_burnin,
                          "acc_rate": acc_rate, "duration": duration, "y_hat": y_hat}

        model_specs = {"reference": self.reference_cell_type, "formula": self.formula}

        return res.CAResultConverter(posterior=posterior,
                                     posterior_predictive=posterior_predictive,
                                     observed_data=observed_data,
                                     dims=dims,
                                     sample_stats=sample_stats,
                                     coords=coords).to_result_data(sampling_stats=sampling_stats,
                                                                   model_specs=model_specs)


class ReferenceModel(CompositionalModel):
    """
    Statistical model for single-cell differential composition analysis with specification of a reference cell type.
    This is the standard scCODA model and recommenced for all uses.

    The hierarchical formulation of the model for one sample is:

    .. math::
         y|x &\\sim DirMult(a(x), \\bar{y}) \\\\
         \\log(a(x)) &= \\alpha + x \\beta \\\\
         \\alpha_k &\\sim N(0, 5) \\quad &\\forall k \\in [K] \\\\
         \\beta_{d, \\hat{k}} &= 0 &\\forall d \\in [D]\\\\
         \\beta_{d, k} &= \\tau_{d, k} \\tilde{\\beta}_{d, k} \\quad &\\forall d \\in [D], k \\in \\{[K] \\smallsetminus \\hat{k}\\} \\\\
         \\tau_{d, k} &= \\frac{\\exp(t_{d, k})}{1+ \\exp(t_{d, k})} \\quad &\\forall d \\in [D], k \\in \\{[K] \\smallsetminus \\hat{k}\\} \\\\
         \\frac{t_{d, k}}{50} &\\sim N(0, 1) \\quad &\\forall d \\in [D], k \\in \\{[K] \\smallsetminus \\hat{k}\\} \\\\
         \\tilde{\\beta}_{d, k} &= (\\tilde{\\mu} + \\tilde{\\sigma}^2) \\cdot \\tilde{\\gamma}_{d, k} \\quad &\\forall d \\in [D], k \\in \\{[K] \\smallsetminus \\hat{k}\\} \\\\
         \\tilde{\\mu} &\\sim N(0, 1) \\\\
         \\tilde{\\sigma}^2 &\\sim HC(0, 1) \\\\
         \\tilde{\\gamma}_{d, k} &\\sim N(0,1) \\quad &\\forall d \\in [D], k \\in \\{[K] \\smallsetminus \\hat{k}\\} \\\\

    with y being the cell counts and x the covariates.

    For further information, see `scCODA: A Bayesian model for compositional single-cell data analysis`
    (Büttner, Ostner et al., 2020)

    """

    def __init__(
            self,
            reference_cell_type: int,
            *args,
            **kwargs):
        """
        Constructor of model class. Defines model structure, log-probability function, parameter names,
        and MCMC starting values.

        Parameters
        ----------
        reference_cell_type
            Index of reference cell type (column in count data matrix)
        args
            arguments passed to top-level class
        kwargs
            arguments passed to top-level class
        """

        super(self.__class__, self).__init__(*args, **kwargs)

        self.reference_cell_type = reference_cell_type
        dtype = tf.float64

        # All parameters that are returned for analysis
        self.param_names = ["mu_b", "sigma_b", "b_offset", "ind_raw", "alpha",
                            "ind", "b_raw", "beta", "concentration", "prediction"]

        alpha_size = [self.K]
        beta_size = [self.D, self.K]
        beta_nobl_size = [self.D, self.K-1]

        Root = tfd.JointDistributionCoroutine.Root

        def model():
            mu_b = yield Root(tfd.Independent(
                tfd.Normal(loc=tf.zeros(1, dtype=dtype),
                           scale=tf.ones(1, dtype=dtype),
                           name="mu_b"),
                reinterpreted_batch_ndims=1))

            sigma_b = yield Root(tfd.Independent(
                tfd.HalfCauchy(tf.zeros(1, dtype=dtype),
                               tf.ones(1, dtype=dtype),
                               name="sigma_b"),
                reinterpreted_batch_ndims=1))

            b_offset = yield Root(tfd.Independent(
                tfd.Normal(
                    loc=tf.zeros(beta_nobl_size, dtype=dtype),
                    scale=tf.ones(beta_nobl_size, dtype=dtype),
                    name="b_offset"),
                reinterpreted_batch_ndims=2))

            # Spike-and-slab
            ind_raw = yield Root(tfd.Independent(
                tfd.Normal(
                    loc=tf.zeros(shape=beta_nobl_size, dtype=dtype),
                    scale=tf.ones(shape=beta_nobl_size, dtype=dtype),
                    name='ind_raw'),
                reinterpreted_batch_ndims=2))

            ind_scaled = ind_raw * 50
            ind = tf.exp(ind_scaled) / (1 + tf.exp(ind_scaled))

            b_raw = mu_b + sigma_b * b_offset

            beta = ind * b_raw

            # Include slope 0 for reference cell type
            beta = tf.concat(axis=1, values=[beta[:, :reference_cell_type],
                                             tf.zeros(shape=[self.D, 1], dtype=dtype),
                                             beta[:, reference_cell_type:]])

            alpha = yield Root(tfd.Independent(
                tfd.Normal(
                    loc=tf.zeros(alpha_size, dtype=dtype),
                    scale=tf.ones(alpha_size, dtype=dtype) * 5,
                    name="alpha"),
                reinterpreted_batch_ndims=1))

            concentrations = tf.exp(alpha + tf.matmul(self.x, beta))

            # Cell count prediction via DirMult
            predictions = yield Root(tfd.Independent(
                tfd.DirichletMultinomial(
                    total_count=tf.cast(self.n_total, dtype),
                    concentration=concentrations,
                    name="predictions"),
                reinterpreted_batch_ndims=1))

        self.model_struct = tfd.JointDistributionCoroutine(model)

        # Joint posterior distribution
        self.target_log_prob_fn = lambda *args:\
            self.model_struct.log_prob(list(args) + [tf.cast(self.y, dtype)])

        # MCMC starting values
        self.init_params = [
            tf.zeros(1, name="init_mu_b", dtype=dtype),
            tf.ones(1, name="init_sigma_b", dtype=dtype),
            tf.random.normal(beta_nobl_size, 0, 1, name='init_b_offset', dtype=dtype),
            tf.zeros(beta_nobl_size, name='init_ind_raw', dtype=dtype),
            tf.random.normal(alpha_size, 0, 1, name='init_alpha', dtype=dtype)
        ]

    # Calculate predicted cell counts (for analysis purposes)
    def get_y_hat(
            self,
            states_burnin: List[any],
            num_results: int,
            num_burnin: int
    ) -> np.ndarray:
        """
        Calculate posterior mode of cell counts (for analysis purposes) and add intermediate parameters
        that are no priors to MCMC results.

        Parameters
        ----------
        states_burnin
            MCMC chain without burn-in samples
        num_results
            Chain length (with burn-in)
        num_burnin
            Number of burn-in samples

        Returns
        -------
        posterior mode

        y_mean
            posterior mode of cell counts
        """

        chain_size_y = [num_results - num_burnin, self.N, self.K]
        chain_size_beta = [num_results - num_burnin, self.D, self.K]

        alphas = states_burnin[4]
        alphas_final = alphas.mean(axis=0)

        ind_raw = states_burnin[3] * 50
        mu_b = states_burnin[0]
        sigma_b = states_burnin[1]
        b_offset = states_burnin[2]

        ind_ = np.exp(ind_raw) / (1 + np.exp(ind_raw))

        b_raw_ = mu_b.reshape((num_results - num_burnin, 1, 1)) + np.einsum("...jk, ...j->...jk", b_offset, sigma_b)

        beta_temp = np.einsum("..., ...", ind_, b_raw_)

        beta_ = np.zeros(chain_size_beta)
        for i in range(num_results - num_burnin):
            beta_[i] = np.concatenate([beta_temp[i, :, :self.reference_cell_type],
                                       np.zeros(shape=[self.D, 1], dtype=np.float64),
                                       beta_temp[i, :, self.reference_cell_type:]], axis=1)
        conc_ = np.exp(np.einsum("jk, ...kl->...jl", self.x, beta_)
                       + alphas.reshape((num_results - num_burnin, 1, self.K)))

        predictions_ = np.zeros(chain_size_y)
        for i in range(num_results - num_burnin):
            pred = tfd.DirichletMultinomial(self.n_total, conc_[i, :, :]).mean().numpy()
            predictions_[i, :, :] = pred

        betas_final = beta_.mean(axis=0)
        states_burnin.append(ind_)
        states_burnin.append(b_raw_)
        states_burnin.append(beta_)
        states_burnin.append(conc_)
        states_burnin.append(predictions_)

        concentration = np.exp(np.matmul(self.x, betas_final) + alphas_final).astype(np.float64)

        y_mean = concentration / np.sum(concentration, axis=1, keepdims=True) * self.n_total.numpy()[:, np.newaxis]

        return y_mean

tensorflow.exp tensorflow.convert_to_tensor tensorflow.zeros tensorflow.random.normal tensorflow.ones tensorflow.matmul tensorflow.cast