# 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/origami.py
# def origami_xl5():
# def origami_l5():
# def origami_l3():
# def origami_m3():
# def two_d_dorfman():
# def origami_m2():
#   def get_groups(self, rng, state):
# class Origami(group_selector.GroupSelector):
#   ORIGAMIS = {
#       271: origami_xl5,
#       68: origami_l5,
#       70: origami_m3,
#       231: origami_m2,
#       274: origami_l3,
#       96 : two_d_dorfman
#   }
# 
# Path: group_testing/group_selectors/random.py
# class RandomSelector(group_selector.GroupSelector):
# class Mezard(group_selector.GroupSelector):
# class RandomDisjoint(group_selector.GroupSelector):
#   def get_groups(self, rng, state):
#   def __init__(self, group_size: int = None):
#   def get_groups(self, rng, state):
# def macula_matrix(d: int, k: int, n: int):
# def sample_groups_of_size(shape: Tuple[int], group_size: int):
# def count_mean(values_arr: onp.ndarray):
# def count_min(values_arr: onp.ndarray):
# def eval_disjoint(mat: onp.ndarray, d: int, count_fn):
# def is_disjoint(mat: onp.ndarray, d: int):
# def sample_disjoint_matrix(num_cols: int,
#                            num_rows: int,
#                            n_max_test: int,
#                            d: int,
#                            max_tries=1e2):
# def sample_maxeval_disjoint_matrix(num_cols: int,
#                                    num_rows: int,
#                                    n_max_test: int,
#                                    d: int,
#                                    max_tries=100,
#                                    count_fn=count_mean):
#   def __init__(self, max_iter=1, method='count', count_fn=count_mean):
#   def get_groups(self, state):
# 
# Path: group_testing/group_selectors/split.py
# class SplitSelector(group_selector.GroupSelector):
# class SplitPositive(group_selector.GroupSelector):
# class TwoDDorfmanPostSelector(group_selector.GroupSelector):
#   def __init__(self, split_factor=None):
#   def get_groups(self, rng, state):
#   def __call__(self, rng, state):
#   def __init__(self, split_factor=2):
#   def _split_groups(self, groups):
#   def get_groups(self, rng, state):
#   def __call__(self, rng, state):
#   def __call__(self, rng, state):
#   NEEDS_POSTERIOR = False


import gin
from group_testing.group_selectors import mutual_information
from group_testing.group_selectors import origami
from group_testing.group_selectors import random
from group_testing.group_selectors import split
class MimaxPolicy(Policy):
  """A policy that starts with Mimax."""

  def __init__(self, subsequent_selectors=None, **kwargs):
    selectors = [mutual_information.MaxMutualInformation(**kwargs)]
    if subsequent_selectors is not None:
      selectors.extend(
          subsequent_selectors if isinstance(subsequent_selectors, list
                                            ) else [subsequent_selectors])
    super().__init__(selectors=selectors)


@gin.configurable
class OrigamiPolicy(Policy):
  """A policy that starts with Origami."""

  def __init__(self, subsequent_selectors=None):
    selectors = [origami.Origami()]
    if subsequent_selectors is not None:
      selectors.extend(
          subsequent_selectors if isinstance(subsequent_selectors, list
                                            ) else [subsequent_selectors])
    super().__init__(selectors=selectors)


@gin.configurable
class MezardPolicy(Policy):
  """A policy that starts with Mezard (random) selector."""

  def __init__(self, subsequent_selectors=None, **kwargs):

# 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/origami.py
# def origami_xl5():
# def origami_l5():
# def origami_l3():
# def origami_m3():
# def two_d_dorfman():
# def origami_m2():
#   def get_groups(self, rng, state):
# class Origami(group_selector.GroupSelector):
#   ORIGAMIS = {
#       271: origami_xl5,
#       68: origami_l5,
#       70: origami_m3,
#       231: origami_m2,
#       274: origami_l3,
#       96 : two_d_dorfman
#   }
# 
# Path: group_testing/group_selectors/random.py
# class RandomSelector(group_selector.GroupSelector):
# class Mezard(group_selector.GroupSelector):
# class RandomDisjoint(group_selector.GroupSelector):
#   def get_groups(self, rng, state):
#   def __init__(self, group_size: int = None):
#   def get_groups(self, rng, state):
# def macula_matrix(d: int, k: int, n: int):
# def sample_groups_of_size(shape: Tuple[int], group_size: int):
# def count_mean(values_arr: onp.ndarray):
# def count_min(values_arr: onp.ndarray):
# def eval_disjoint(mat: onp.ndarray, d: int, count_fn):
# def is_disjoint(mat: onp.ndarray, d: int):
# def sample_disjoint_matrix(num_cols: int,
#                            num_rows: int,
#                            n_max_test: int,
#                            d: int,
#                            max_tries=1e2):
# def sample_maxeval_disjoint_matrix(num_cols: int,
#                                    num_rows: int,
#                                    n_max_test: int,
#                                    d: int,
#                                    max_tries=100,
#                                    count_fn=count_mean):
#   def __init__(self, max_iter=1, method='count', count_fn=count_mean):
#   def get_groups(self, state):
# 
# Path: group_testing/group_selectors/split.py
# class SplitSelector(group_selector.GroupSelector):
# class SplitPositive(group_selector.GroupSelector):
# class TwoDDorfmanPostSelector(group_selector.GroupSelector):
#   def __init__(self, split_factor=None):
#   def get_groups(self, rng, state):
#   def __call__(self, rng, state):
#   def __init__(self, split_factor=2):
#   def _split_groups(self, groups):
#   def get_groups(self, rng, state):
#   def __call__(self, rng, state):
#   def __call__(self, rng, state):
#   NEEDS_POSTERIOR = False


# Path: group_testing/policy.py
import gin
from group_testing.group_selectors import mutual_information
from group_testing.group_selectors import origami
from group_testing.group_selectors import random
from group_testing.group_selectors import split

class MimaxPolicy(Policy):
  """A policy that starts with Mimax."""

  def __init__(self, subsequent_selectors=None, **kwargs):
    selectors = [mutual_information.MaxMutualInformation(**kwargs)]
    if subsequent_selectors is not None:
      selectors.extend(
          subsequent_selectors if isinstance(subsequent_selectors, list
                                            ) else [subsequent_selectors])
    super().__init__(selectors=selectors)


@gin.configurable
class OrigamiPolicy(Policy):
  """A policy that starts with Origami."""

  def __init__(self, subsequent_selectors=None):
    selectors = [origami.Origami()]
    if subsequent_selectors is not None:
      selectors.extend(
          subsequent_selectors if isinstance(subsequent_selectors, list
                                            ) else [subsequent_selectors])
    super().__init__(selectors=selectors)


@gin.configurable
class MezardPolicy(Policy):
  """A policy that starts with Mezard (random) selector."""

  def __init__(self, subsequent_selectors=None, **kwargs):

    selectors = [random.Mezard(**kwargs)]