# 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:


import csv
import gin
import jax
import numpy as np
from typing import Union, Iterable
from group_testing import utils
     base_infection_rate: Can be either a single value between 0 and 1, which
      applies to all patients, or a prior value for each patient given as a list
      or array. In the case one wants to enforce a certain infection status
      (not a random one), it suffices to set these values to 0 and 1 for non
      infected and infected patients respectively.
     specificity: the specificity of the PCR machine to be used. In case of an
      vector, it corresponds to the specificity for a given group size.
     sensitivity: the sensitivity of the PCR machine to be used. In case of an
      vector, it corresponds to the sensitivity for a given group size.
    """
    if num_patients < 2:
      raise ValueError('Number of patients should be larger than 2')

    self.num_patients = num_patients
    self.freeze_diseased = freeze_diseased
    self._base_infection_rate = np.atleast_1d(base_infection_rate)
    self._specificity = np.atleast_1d(specificity)
    self._sensitivity = np.atleast_1d(sensitivity)
    self.diseased = None

  def reset(self, rng):
    if not self.freeze_diseased or self.diseased is None:
      random_matrix = jax.random.uniform(rng, (self.num_patients,))
      self.diseased = random_matrix < self._base_infection_rate

  def group_tests_outputs(self, rng, groups):
    """Produces test outputs taking into account test errors."""
    n_groups = groups.shape[0]
    group_disease_indicator = np.dot(groups, self.diseased) > 0
    group_sizes = np.sum(groups, axis=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/wet_lab.py
import csv
import gin
import jax
import numpy as np
from typing import Union, Iterable
from group_testing import utils

     base_infection_rate: Can be either a single value between 0 and 1, which
      applies to all patients, or a prior value for each patient given as a list
      or array. In the case one wants to enforce a certain infection status
      (not a random one), it suffices to set these values to 0 and 1 for non
      infected and infected patients respectively.
     specificity: the specificity of the PCR machine to be used. In case of an
      vector, it corresponds to the specificity for a given group size.
     sensitivity: the sensitivity of the PCR machine to be used. In case of an
      vector, it corresponds to the sensitivity for a given group size.
    """
    if num_patients < 2:
      raise ValueError('Number of patients should be larger than 2')

    self.num_patients = num_patients
    self.freeze_diseased = freeze_diseased
    self._base_infection_rate = np.atleast_1d(base_infection_rate)
    self._specificity = np.atleast_1d(specificity)
    self._sensitivity = np.atleast_1d(sensitivity)
    self.diseased = None

  def reset(self, rng):
    if not self.freeze_diseased or self.diseased is None:
      random_matrix = jax.random.uniform(rng, (self.num_patients,))
      self.diseased = random_matrix < self._base_infection_rate

  def group_tests_outputs(self, rng, groups):
    """Produces test outputs taking into account test errors."""
    n_groups = groups.shape[0]
    group_disease_indicator = np.dot(groups, self.diseased) > 0
    group_sizes = np.sum(groups, axis=1)

    specificity = utils.select_from_sizes(self._specificity, group_sizes)