# 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 itertools
import gin
import jax
import jax.numpy as np
import numpy as onp
import scipy.special
from typing import Tuple
from group_testing import utils
from group_testing.group_selectors import group_selector
  """Selects groups randomly."""

  def get_groups(self, rng, state):
    shape = (state.extra_tests_needed, state.num_patients)
    threshold = state.max_group_size / state.num_patients
    return jax.random.uniform(rng, shape=shape) < threshold


@gin.configurable
class Mezard(group_selector.GroupSelector):
  """Selects groups randomly with predefined/adaptive-to-infection-rate size."""

  def __init__(self, group_size: int = None):
    super().__init__()
    self.group_size = group_size

  def get_groups(self, rng, state):
    """Produces random design matrix fixed number of 1s per line.

    Args:
     rng: np.ndarray<int>[2]: the random key.
     state: the current state.State of the system.

    Returns:
     A np.array<bool>[num_groups, patients].
    """
    if self.group_size is None:
      if np.size(state.prior_infection_rate) == 1:
        # candidate group sizes
        group_sizes = np.arange(state.max_group_size) + 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/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/random.py
import itertools
import gin
import jax
import jax.numpy as np
import numpy as onp
import scipy.special
from typing import Tuple
from group_testing import utils
from group_testing.group_selectors import group_selector

  """Selects groups randomly."""

  def get_groups(self, rng, state):
    shape = (state.extra_tests_needed, state.num_patients)
    threshold = state.max_group_size / state.num_patients
    return jax.random.uniform(rng, shape=shape) < threshold


@gin.configurable
class Mezard(group_selector.GroupSelector):
  """Selects groups randomly with predefined/adaptive-to-infection-rate size."""

  def __init__(self, group_size: int = None):
    super().__init__()
    self.group_size = group_size

  def get_groups(self, rng, state):
    """Produces random design matrix fixed number of 1s per line.

    Args:
     rng: np.ndarray<int>[2]: the random key.
     state: the current state.State of the system.

    Returns:
     A np.array<bool>[num_groups, patients].
    """
    if self.group_size is None:
      if np.size(state.prior_infection_rate) == 1:
        # candidate group sizes
        group_sizes = np.arange(state.max_group_size) + 1

        sensitivity = utils.select_from_sizes(state.prior_sensitivity,