# Path: group_testing/bayes.py
# def log_likelihood(particles: np.ndarray, test_results: np.ndarray,
#                    groups: np.ndarray, log_specificity: np.ndarray,
#                    log_1msensitivity: np.ndarray) -> float:
# def log_prior(particles: np.ndarray,
#               base_infection_rate: np.ndarray) -> np.ndarray:
# def log_probability(particles: np.ndarray,
#                     test_results: np.ndarray,
#                     groups: np.ndarray,
#                     log_prior_specificity: np.ndarray,
#                     log_prior_1msensitivity: np.ndarray,
#                     prior_infection_rate: np.ndarray):
# def tempered_logpos_logbase(particles: np.ndarray,
#                             log_posterior_params: Dict[str, np.ndarray],
#                             log_base_measure_params: Dict[str, np.ndarray],
#                             temperature: float):


import collections
import gin
import jax
import jax.nn
import jax.numpy as np
import jax.scipy.special
from typing import Dict
from group_testing import bayes
  The Liu variant comes from Jun Liu's remarks in
  https://academic.oup.com/biomet/article-abstract/83/3/681/241540?redirectedFrom=fulltext

  which essentially changes the acceptance of a flip from
    p(flip) / [ p(no flip) + p(flip) ]
  to
    min(1, p(flip) / p(no flip) )

  In other words, Liu's modification increases the probability to flip.

  Args:
   rng: np.ndarray<int> random key.
   particles: np.ndarray [n_particles,n_patients] plausible infections states.
   rho: float, scaling for posterior.
   log_posterior_params: Dict of parameters to compute log-posterior.
   log_base_measure_params: Dict of parameters to compute log-base measure.
   cycles: the number of times we want of do Gibbs sampling.
   liu_modification : use or not Liu's modification.

  Returns:
   A np.array representing the new particles.
  """

  def gibbs_loop(i, rng_particles_log_posteriors):
    rng, particles, log_posteriors = rng_particles_log_posteriors
    i = i % num_patients
    # flip values at index i
    particles_flipped = jax.ops.index_update(particles, jax.ops.index[:, i],
                                             np.logical_not(particles[:, i]))
    # compute log_posterior of flipped particles

# Path: group_testing/bayes.py
# def log_likelihood(particles: np.ndarray, test_results: np.ndarray,
#                    groups: np.ndarray, log_specificity: np.ndarray,
#                    log_1msensitivity: np.ndarray) -> float:
# def log_prior(particles: np.ndarray,
#               base_infection_rate: np.ndarray) -> np.ndarray:
# def log_probability(particles: np.ndarray,
#                     test_results: np.ndarray,
#                     groups: np.ndarray,
#                     log_prior_specificity: np.ndarray,
#                     log_prior_1msensitivity: np.ndarray,
#                     prior_infection_rate: np.ndarray):
# def tempered_logpos_logbase(particles: np.ndarray,
#                             log_posterior_params: Dict[str, np.ndarray],
#                             log_base_measure_params: Dict[str, np.ndarray],
#                             temperature: float):


# Path: group_testing/samplers/kernels.py
import collections
import gin
import jax
import jax.nn
import jax.numpy as np
import jax.scipy.special
from typing import Dict
from group_testing import bayes

  The Liu variant comes from Jun Liu's remarks in
  https://academic.oup.com/biomet/article-abstract/83/3/681/241540?redirectedFrom=fulltext

  which essentially changes the acceptance of a flip from
    p(flip) / [ p(no flip) + p(flip) ]
  to
    min(1, p(flip) / p(no flip) )

  In other words, Liu's modification increases the probability to flip.

  Args:
   rng: np.ndarray<int> random key.
   particles: np.ndarray [n_particles,n_patients] plausible infections states.
   rho: float, scaling for posterior.
   log_posterior_params: Dict of parameters to compute log-posterior.
   log_base_measure_params: Dict of parameters to compute log-base measure.
   cycles: the number of times we want of do Gibbs sampling.
   liu_modification : use or not Liu's modification.

  Returns:
   A np.array representing the new particles.
  """

  def gibbs_loop(i, rng_particles_log_posteriors):
    rng, particles, log_posteriors = rng_particles_log_posteriors
    i = i % num_patients
    # flip values at index i
    particles_flipped = jax.ops.index_update(particles, jax.ops.index[:, i],
                                             np.logical_not(particles[:, i]))
    # compute log_posterior of flipped particles

    log_posteriors_flipped_at_i = bayes.tempered_logpos_logbase(