# 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
    cur_group = candidate_groups[index, :]

  else:
    # return most promising group by adding a 1
    cur_group = jax.ops.index_update(
        cur_group, indices_of_false_in_cur_group[index], True)

  # refresh the status of vector positives
  cur_positives = positive_in_groups[index, :]
  new_objective = objectives[index]
  prob_particles_states = upd_prob_groups_x_particle_x_states(
      prob_prev_groups_of_particles_x_states, cur_positives,
      specificity, rho)
  new_cond_entropy = cond_entropy[index]
  return (cur_group, cur_positives, new_objective,
          prob_particles_states, new_cond_entropy)


def add_ones_to_line(single_group):
  indices, = np.where(np.logical_not(single_group))
  k_zeros = len(indices)
  groups = (
      np.zeros((k_zeros, single_group.shape[0]), dtype=bool)
      + single_group[np.newaxis, :])
  groups = jax.ops.index_update(
      groups, jax.ops.index[np.arange(0, k_zeros), indices], True)
  return groups


@gin.configurable

# 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

    cur_group = candidate_groups[index, :]

  else:
    # return most promising group by adding a 1
    cur_group = jax.ops.index_update(
        cur_group, indices_of_false_in_cur_group[index], True)

  # refresh the status of vector positives
  cur_positives = positive_in_groups[index, :]
  new_objective = objectives[index]
  prob_particles_states = upd_prob_groups_x_particle_x_states(
      prob_prev_groups_of_particles_x_states, cur_positives,
      specificity, rho)
  new_cond_entropy = cond_entropy[index]
  return (cur_group, cur_positives, new_objective,
          prob_particles_states, new_cond_entropy)


def add_ones_to_line(single_group):
  indices, = np.where(np.logical_not(single_group))
  k_zeros = len(indices)
  groups = (
      np.zeros((k_zeros, single_group.shape[0]), dtype=bool)
      + single_group[np.newaxis, :])
  groups = jax.ops.index_update(
      groups, jax.ops.index[np.arange(0, k_zeros), indices], True)
  return groups


@gin.configurable

class MaxMutualInformation(group_selector.GroupSelector):