# Path: group_testing/metrics.py
# def auc(y_true, y_pred):
# def kl(x, y):
# def binary_entropy(u):
# def entropy(table, axis=0):
#   def __init__(self,
#                workdir=None,
#                num_simulations=0,
#                num_cycles=0,
#                num_patients=0,
#                num_tests_per_cycle=0):
#   def update(self,
#              simulation,
#              cycle,
#              marginal,
#              ground_truth,
#              groups=None,
#              test_results=None,
#              lbp_convergence=None,
#              smc_convergence=None):
#   def load(self):
#   def export(self):
#   def _merge_one(self, other, name: str):
#   def merge(self, other):
#   def _extract_data(self, simulation=None, cycle: int = -1):
#   def roc_curve(self, simulation=None, cycle=-1):
#   def precision_recall_curve(self, simulation: int = None, cycle=-1):
# class Metrics:
# 
# 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/group_selectors/group_selector.py
# class GroupSelector:
#   NEEDS_POSTERIOR = False
#   def __call__(self, rng, state):
#   def get_groups(self, rng, state):


import functools
import gin
import jax
import jax.numpy as np
from absl import logging
from group_testing import metrics
from group_testing import utils
from group_testing.group_selectors import group_selector
@jax.jit
@functools.partial(jax.vmap, in_axes=[0, None, None, None, None])
def compute_whole_entropy_groups(positive_in_group,
                                 specificity,
                                 rho,
                                 particle_weights,
                                 prob_prev_groups_of_particles_x_states):
  """For a j-th candidate group, compute entropy of 2^{j} possible outcomes.

  Args:
    positive_in_group: [num_particles,] array of bools, whether group of
      interest has a positive patient in the particle considered.
    specificity: float, specificity of test for group size under consideration
    rho: float, combination of specificity and sensitivity.
    particle_weights: [num_particles,] array of floats, sum to 1.
    prob_prev_groups_of_particles_x_states: [num_particles, 2^k] prob. table  of
      k tests taking one of 2^k possible configurations when ground state is
      assumed to be one of the particles.
  Returns:
    a float
  """

  prob_new_test = probabilities_new_test(positive_in_group, specificity, rho)
  # we now incorporate previous probability of all previous groups added so far
  # and expand x 2 the state space of possible test results, and sum to get
  # exactly the 2 x 2^k probabilities of each combination of k+1 test results
  n_p_prev_groups_prob_states = np.dot(
      prob_prev_groups_of_particles_x_states.T,
      prob_new_test * particle_weights[:, np.newaxis])
  # compute entropy of this 2 x 2^k matrix of probabilities.

# Path: group_testing/metrics.py
# def auc(y_true, y_pred):
# def kl(x, y):
# def binary_entropy(u):
# def entropy(table, axis=0):
#   def __init__(self,
#                workdir=None,
#                num_simulations=0,
#                num_cycles=0,
#                num_patients=0,
#                num_tests_per_cycle=0):
#   def update(self,
#              simulation,
#              cycle,
#              marginal,
#              ground_truth,
#              groups=None,
#              test_results=None,
#              lbp_convergence=None,
#              smc_convergence=None):
#   def load(self):
#   def export(self):
#   def _merge_one(self, other, name: str):
#   def merge(self, other):
#   def _extract_data(self, simulation=None, cycle: int = -1):
#   def roc_curve(self, simulation=None, cycle=-1):
#   def precision_recall_curve(self, simulation: int = None, cycle=-1):
# class Metrics:
# 
# 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/group_selectors/group_selector.py
# class GroupSelector:
#   NEEDS_POSTERIOR = False
#   def __call__(self, rng, state):
#   def get_groups(self, rng, state):


# Path: group_testing/group_selectors/mutual_information.py
import functools
import gin
import jax
import jax.numpy as np
from absl import logging
from group_testing import metrics
from group_testing import utils
from group_testing.group_selectors import group_selector

@jax.jit
@functools.partial(jax.vmap, in_axes=[0, None, None, None, None])
def compute_whole_entropy_groups(positive_in_group,
                                 specificity,
                                 rho,
                                 particle_weights,
                                 prob_prev_groups_of_particles_x_states):
  """For a j-th candidate group, compute entropy of 2^{j} possible outcomes.

  Args:
    positive_in_group: [num_particles,] array of bools, whether group of
      interest has a positive patient in the particle considered.
    specificity: float, specificity of test for group size under consideration
    rho: float, combination of specificity and sensitivity.
    particle_weights: [num_particles,] array of floats, sum to 1.
    prob_prev_groups_of_particles_x_states: [num_particles, 2^k] prob. table  of
      k tests taking one of 2^k possible configurations when ground state is
      assumed to be one of the particles.
  Returns:
    a float
  """

  prob_new_test = probabilities_new_test(positive_in_group, specificity, rho)
  # we now incorporate previous probability of all previous groups added so far
  # and expand x 2 the state space of possible test results, and sum to get
  # exactly the 2 x 2^k probabilities of each combination of k+1 test results
  n_p_prev_groups_prob_states = np.dot(
      prob_prev_groups_of_particles_x_states.T,
      prob_new_test * particle_weights[:, np.newaxis])
  # compute entropy of this 2 x 2^k matrix of probabilities.

  return metrics.entropy(n_p_prev_groups_prob_states, axis=None)