# Path: group_testing/bayes.py
# def log_likelihood(particles: np.ndarray, test_results: np.ndarray,
#                    groups: np.ndarray, log_specificity: np.ndarray,
#                    log_1msensitivity: np.ndarray) -> float:
# def log_prior(particles: np.ndarray,
#               base_infection_rate: np.ndarray) -> np.ndarray:
# def log_probability(particles: np.ndarray,
#                     test_results: np.ndarray,
#                     groups: np.ndarray,
#                     log_prior_specificity: np.ndarray,
#                     log_prior_1msensitivity: np.ndarray,
#                     prior_infection_rate: np.ndarray):
# def tempered_logpos_logbase(particles: np.ndarray,
#                             log_posterior_params: Dict[str, np.ndarray],
#                             log_base_measure_params: Dict[str, np.ndarray],
#                             temperature: float):
# 
# Path: group_testing/samplers/sampler.py
# class Sampler:
#   def __init__(self):
#   def reset(self):
#   def is_cheap(self):
#   def marginal(self):
#   def reset_convergence_metric(self):
# 
# Path: group_testing/samplers/temperature.py
# def effective_sample_size(alpha: float,
#                           log_posterior: np.ndarray):
# def find_step_length(rho: float,
#                      log_posterior: np.ndarray,
#                      tolerance=0.02,
#                      effective_sample_size_target=0.9):
# def importance_weights(log_unnormalized_probabilities):


import gin
import jax.numpy as np
from group_testing import bayes
from group_testing.samplers import sampler
from group_testing.samplers import temperature
def all_binary_vectors(num_samples, upper_bound: float = 1.0):
  """Generates all 2^n binary vectors."""
  # TODO(cuturi) use itertools ?
  results = np.array([[0], [1]], dtype=bool)
  bound_per_line = np.ceil(upper_bound * num_samples) + 1
  for _ in range(num_samples - 1):
    num_rows = results.shape[0]
    twice_binary_vectors = np.concatenate((results, results), axis=0)
    zo = np.concatenate((np.zeros((num_rows, 1), dtype=bool),
                         np.ones((num_rows, 1), dtype=bool)),
                        axis=0)
    results = np.concatenate((twice_binary_vectors, zo), axis=1)
    if 0 < upper_bound < 1.0:
      count = np.sum(results, axis=-1)
      results = results[(count < bound_per_line), :]
  return results


@gin.configurable
class ExhaustiveSampler(sampler.Sampler):
  """Produces all possible binary vectors (with up to upper bound positives)."""

  NAME = 'Exh'

  def __init__(self, upper_bound: float = 0.5):
    super().__init__()
    self.upper_bound = upper_bound

  def produce_sample(self, rng, state):
    self.particles = all_binary_vectors(state.num_patients, self.upper_bound)

# Path: group_testing/bayes.py
# def log_likelihood(particles: np.ndarray, test_results: np.ndarray,
#                    groups: np.ndarray, log_specificity: np.ndarray,
#                    log_1msensitivity: np.ndarray) -> float:
# def log_prior(particles: np.ndarray,
#               base_infection_rate: np.ndarray) -> np.ndarray:
# def log_probability(particles: np.ndarray,
#                     test_results: np.ndarray,
#                     groups: np.ndarray,
#                     log_prior_specificity: np.ndarray,
#                     log_prior_1msensitivity: np.ndarray,
#                     prior_infection_rate: np.ndarray):
# def tempered_logpos_logbase(particles: np.ndarray,
#                             log_posterior_params: Dict[str, np.ndarray],
#                             log_base_measure_params: Dict[str, np.ndarray],
#                             temperature: float):
# 
# Path: group_testing/samplers/sampler.py
# class Sampler:
#   def __init__(self):
#   def reset(self):
#   def is_cheap(self):
#   def marginal(self):
#   def reset_convergence_metric(self):
# 
# Path: group_testing/samplers/temperature.py
# def effective_sample_size(alpha: float,
#                           log_posterior: np.ndarray):
# def find_step_length(rho: float,
#                      log_posterior: np.ndarray,
#                      tolerance=0.02,
#                      effective_sample_size_target=0.9):
# def importance_weights(log_unnormalized_probabilities):


# Path: group_testing/samplers/exhaustive.py
import gin
import jax.numpy as np
from group_testing import bayes
from group_testing.samplers import sampler
from group_testing.samplers import temperature

def all_binary_vectors(num_samples, upper_bound: float = 1.0):
  """Generates all 2^n binary vectors."""
  # TODO(cuturi) use itertools ?
  results = np.array([[0], [1]], dtype=bool)
  bound_per_line = np.ceil(upper_bound * num_samples) + 1
  for _ in range(num_samples - 1):
    num_rows = results.shape[0]
    twice_binary_vectors = np.concatenate((results, results), axis=0)
    zo = np.concatenate((np.zeros((num_rows, 1), dtype=bool),
                         np.ones((num_rows, 1), dtype=bool)),
                        axis=0)
    results = np.concatenate((twice_binary_vectors, zo), axis=1)
    if 0 < upper_bound < 1.0:
      count = np.sum(results, axis=-1)
      results = results[(count < bound_per_line), :]
  return results


@gin.configurable
class ExhaustiveSampler(sampler.Sampler):
  """Produces all possible binary vectors (with up to upper bound positives)."""

  NAME = 'Exh'

  def __init__(self, upper_bound: float = 0.5):
    super().__init__()
    self.upper_bound = upper_bound

  def produce_sample(self, rng, state):
    self.particles = all_binary_vectors(state.num_patients, self.upper_bound)

    log_posteriors = bayes.log_probability(