# Path: group_testing/utils.py
# def unique(rng, binary_vectors: np.ndarray) -> int:
# def select_from_sizes(values: np.ndarray, sizes: np.ndarray) -> np.ndarray:


from typing import Dict
from group_testing import utils
from jax.scipy import special
import jax
import jax.numpy as np
# Lint as: python3
"""Computes the prior, posterior, and likelihood."""



@jax.jit
def log_likelihood(particles: np.ndarray, test_results: np.ndarray,
                   groups: np.ndarray, log_specificity: np.ndarray,
                   log_1msensitivity: np.ndarray) -> float:
  """Computes individual (parallel) log_likelihood of k_groups test results.

  Args:
   particles: np.ndarray<bool>[n_particles, n_patients]. Each one is a possible
     scenario of a disease status of n patients.
   test_results: np.ndarray<bool>[n_groups] the results given by the wet lab for
     each of the tested groups.
   groups: np.ndarray<bool>[num_groups, num_patients] the definition of the
     group that were tested.
   log_specificity: np.ndarray. Depending on the configuration, it can be an
     array of size one or more if we have different sensitivities per group
     size.
   log_1msensitivity: np.ndarray. Depending on the configuration, it can be an
     array of size one or more if we have different specificities per group
     size.

  Returns:
   The log likelihood of the particles given the test results.
  """
  positive_in_groups = np.dot(groups, np.transpose(particles)) > 0
  group_sizes = np.sum(groups, axis=1)

# Path: group_testing/utils.py
# def unique(rng, binary_vectors: np.ndarray) -> int:
# def select_from_sizes(values: np.ndarray, sizes: np.ndarray) -> np.ndarray:


# Path: group_testing/bayes.py
from typing import Dict
from group_testing import utils
from jax.scipy import special
import jax
import jax.numpy as np

# Lint as: python3
"""Computes the prior, posterior, and likelihood."""



@jax.jit
def log_likelihood(particles: np.ndarray, test_results: np.ndarray,
                   groups: np.ndarray, log_specificity: np.ndarray,
                   log_1msensitivity: np.ndarray) -> float:
  """Computes individual (parallel) log_likelihood of k_groups test results.

  Args:
   particles: np.ndarray<bool>[n_particles, n_patients]. Each one is a possible
     scenario of a disease status of n patients.
   test_results: np.ndarray<bool>[n_groups] the results given by the wet lab for
     each of the tested groups.
   groups: np.ndarray<bool>[num_groups, num_patients] the definition of the
     group that were tested.
   log_specificity: np.ndarray. Depending on the configuration, it can be an
     array of size one or more if we have different sensitivities per group
     size.
   log_1msensitivity: np.ndarray. Depending on the configuration, it can be an
     array of size one or more if we have different specificities per group
     size.

  Returns:
   The log likelihood of the particles given the test results.
  """
  positive_in_groups = np.dot(groups, np.transpose(particles)) > 0
  group_sizes = np.sum(groups, axis=1)

  log_specificity = utils.select_from_sizes(log_specificity, group_sizes)