# 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
# def collapse_particles(rng, particle_weights, particles):
# def probabilities_new_test(positive_in_group,
#                            specificity,
#                            rho):
# def upd_prob_groups_x_particle_x_states(
#     prob_prev_groups_of_particles_x_states, positive_in_group,
#     specificity, rho):
# def compute_whole_entropy_groups(positive_in_group,
#                                  specificity,
#                                  rho,
#                                  particle_weights,
#                                  prob_prev_groups_of_particles_x_states):
# def joint_mi_criterion_mg(particle_weights, particles, cur_group, cur_positives,
#                           prob_prev_groups_of_particles_x_states,
#                           previous_groups_cumcond_entropy,
#                           sensitivity,
#                           specificity,
#                           backtracking: bool):
# def add_ones_to_line(single_group):
#   def __init__(self, forward_iterations=1, backward_iterations=0):
#   def get_groups(self, rng, state):
# class MaxMutualInformation(group_selector.GroupSelector):
#   NEEDS_POSTERIOR = True


import itertools
import gin
import jax
import jax.numpy as np
from group_testing import metrics
from group_testing import utils
from group_testing.group_selectors import group_selector
from group_testing.group_selectors import mutual_information
    # Backward mode: test groups obtained by removing an item to cur_group
    candidate_groups = np.logical_not(
        mutual_information.add_ones_to_line(np.logical_not(cur_group)))
  else:
    # Forward mode: test groups obtained by adding an item to cur_group
    candidate_groups = mutual_information.add_ones_to_line(cur_group)
  n_candidates = candidate_groups.shape[0]

  # Combine past groups with candidate groups
  candidate_sets = np.concatenate(
      (np.repeat(previous_groups[:, :, np.newaxis], n_candidates, axis=2),
       np.expand_dims(np.transpose(candidate_groups), axis=0)),
      axis=0)

  # Compute utility of each candidate group
  group_sizes = np.sum(candidate_sets[:, :, 0], axis=1)
  group_sensitivities = utils.select_from_sizes(sensitivity, group_sizes)
  group_specificities = utils.select_from_sizes(specificity, group_sizes)
  group_util_fun = lambda x: group_utility(particle_weights, particles, x,
                                           group_sensitivities,
                                           group_specificities, utility_fun)
  mgroup_util_fun = jax.vmap(group_util_fun, in_axes=2)
  objectives = mgroup_util_fun(candidate_sets)

  # Greedy selection of largest value
  index = np.argmax(objectives)
  return (candidate_groups[index, :], objectives[index])


@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
# def collapse_particles(rng, particle_weights, particles):
# def probabilities_new_test(positive_in_group,
#                            specificity,
#                            rho):
# def upd_prob_groups_x_particle_x_states(
#     prob_prev_groups_of_particles_x_states, positive_in_group,
#     specificity, rho):
# def compute_whole_entropy_groups(positive_in_group,
#                                  specificity,
#                                  rho,
#                                  particle_weights,
#                                  prob_prev_groups_of_particles_x_states):
# def joint_mi_criterion_mg(particle_weights, particles, cur_group, cur_positives,
#                           prob_prev_groups_of_particles_x_states,
#                           previous_groups_cumcond_entropy,
#                           sensitivity,
#                           specificity,
#                           backtracking: bool):
# def add_ones_to_line(single_group):
#   def __init__(self, forward_iterations=1, backward_iterations=0):
#   def get_groups(self, rng, state):
# class MaxMutualInformation(group_selector.GroupSelector):
#   NEEDS_POSTERIOR = True


# Path: group_testing/group_selectors/bayes_oed.py
import itertools
import gin
import jax
import jax.numpy as np
from group_testing import metrics
from group_testing import utils
from group_testing.group_selectors import group_selector
from group_testing.group_selectors import mutual_information

    # Backward mode: test groups obtained by removing an item to cur_group
    candidate_groups = np.logical_not(
        mutual_information.add_ones_to_line(np.logical_not(cur_group)))
  else:
    # Forward mode: test groups obtained by adding an item to cur_group
    candidate_groups = mutual_information.add_ones_to_line(cur_group)
  n_candidates = candidate_groups.shape[0]

  # Combine past groups with candidate groups
  candidate_sets = np.concatenate(
      (np.repeat(previous_groups[:, :, np.newaxis], n_candidates, axis=2),
       np.expand_dims(np.transpose(candidate_groups), axis=0)),
      axis=0)

  # Compute utility of each candidate group
  group_sizes = np.sum(candidate_sets[:, :, 0], axis=1)
  group_sensitivities = utils.select_from_sizes(sensitivity, group_sizes)
  group_specificities = utils.select_from_sizes(specificity, group_sizes)
  group_util_fun = lambda x: group_utility(particle_weights, particles, x,
                                           group_sensitivities,
                                           group_specificities, utility_fun)
  mgroup_util_fun = jax.vmap(group_util_fun, in_axes=2)
  objectives = mgroup_util_fun(candidate_sets)

  # Greedy selection of largest value
  index = np.argmax(objectives)
  return (candidate_groups[index, :], objectives[index])


@gin.configurable

class BayesOED(group_selector.GroupSelector):