# Path: group_testing/policy.py
# class Policy(object):
# class MimaxPolicy(Policy):
# class OrigamiPolicy(Policy):
# class MezardPolicy(Policy):
# class Dorfman(Policy):
# class InformativeDorfmanPolicy(Policy):
#   def __init__(self, selectors=None):
#   def reset(self):
#   def get_selector(self, index=None):
#   def next_selector(self):
#   def act(self, rng, state):
#   def __init__(self, subsequent_selectors=None, **kwargs):
#   def __init__(self, subsequent_selectors=None):
#   def __init__(self, subsequent_selectors=None, **kwargs):
#   def __init__(self, second_split=None, **kwargs):
#   def __init__(self, **kwargs):
# 
# Path: group_testing/simulator.py
# class Simulator(object):
#   def __init__(self,
#                workdir: str = None,
#                wetlab=wet_lab.WetLab(),
#                policy=group_policy.Dorfman(),
#                sampler=smc.SmcSampler(),
#                cheap_sampler=lbp.LbpSampler(),
#                num_simulations: int = 1,
#                num_tests_per_cycle: int = 10,
#                max_test_cycles: int = 5,
#                max_group_size: int = 8,
#                prior_specificity: Union[float, Iterable[float]] = 0.97,
#                prior_sensitivity: Union[float, Iterable[float]] = 0.85,
#                prior_infection_rate: Union[float, Iterable[float]] = 0.05,
#                metrics_cls=metrics.Metrics,
#                export_metrics_every: int = 5):
#   def reset(self, rng: np.ndarray):
#   def sampler(self):
#   def run(self, rngkey=None):
#   def is_over(self) -> bool:
#   def run_simulation(self, rng, sim_idx: int):
#   def run_one_cycle(self, rng, sim_idx: int):
#   def marginal(self):
#   def process_tests_results(self, rng, new_tests):
#   def get_new_groups(self, rng) -> np.ndarray:
#   def _resample(self, rng):
#   def _on_iteration_end(self,
#                         sim_idx: int,
#                         cycle_idx: int,
#                         new_groups=None,
#                         new_tests=None):
#   def interactive_loop(self, rngkey: int = None, show_fn=None):
# 
# Path: group_testing/wet_lab.py
# class WetLab:
# class WetLabCSV(WetLab):
#   def __init__(self,
#                num_patients: int = 2,
#                freeze_diseased: bool = True,
#                base_infection_rate: Union[float, Iterable[float]] = 0.05,
#                specificity: Union[float, Iterable[float]] = 0.97,
#                sensitivity: Union[float, Iterable[float]] = 0.85):
#   def reset(self, rng):
#   def group_tests_outputs(self, rng, groups):
#   def __init__(self, filename: str = None):
#   def reset(self, rng):
#   def group_tests_outputs(self, rng, groups):
#   def to_csv(self, groups: np.ndarray, marginals: np.ndarray):
#     def to_list(arr):
# 
# Path: group_testing/samplers/sequential_monte_carlo.py
# class SmcSampler(sampler.Sampler):
#   NAME = 'smc'
#   def __init__(self,
#                kernel=kernels.Gibbs(),
#                resample_at_each_iteration: bool = False,
#                start_from_prior: bool = True,
#                num_particles: int = 10000,
#                min_kernel_iterations: int = 2,
#                max_kernel_iterations: int = 20,
#                min_ratio_delta: float = 0.02,
#                target_unique_ratio: float = 0.95):
#   def reset(self):
#   def is_cheap(self):
#   def produce_sample(self, rng, state):
#   def resample_move(self,
#                     rng,
#                     particles,
#                     state,
#                     sampling_from_scratch):
#   def resample(self,
#                rng: int,
#                particle_weights: np.ndarray) -> np.ndarray:
#   def move(self,
#            rng: int,
#            particles: np.ndarray,
#            rho: float,
#            log_posterior_params: Dict[str, np.ndarray],
#            log_base_measure_params: Dict[str, np.ndarray]):


import time
import jax
import numpy as np
from group_testing import policy as group_policy
from group_testing import simulator
from group_testing import wet_lab
from group_testing.samplers import sequential_monte_carlo as smc
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# Lint as: python3
"""An interactive loop to go back and forth from colab to spreadsheet."""





class InteractiveSimulator(simulator.Simulator):
  """Interactive group testing loop, reading an write into a spreadsheet."""

  def __init__(self, sheet):
    wetlab = wet_lab.WetLab(num_patients=sheet.num_patients)
    sampler = smc.SmcSampler(num_particles=sheet.num_patients * 150,
                             resample_at_each_iteration=True)

    policy_name = sheet.params['policy']
    if policy_name == 'G-MIMAX':

# Path: group_testing/policy.py
# class Policy(object):
# class MimaxPolicy(Policy):
# class OrigamiPolicy(Policy):
# class MezardPolicy(Policy):
# class Dorfman(Policy):
# class InformativeDorfmanPolicy(Policy):
#   def __init__(self, selectors=None):
#   def reset(self):
#   def get_selector(self, index=None):
#   def next_selector(self):
#   def act(self, rng, state):
#   def __init__(self, subsequent_selectors=None, **kwargs):
#   def __init__(self, subsequent_selectors=None):
#   def __init__(self, subsequent_selectors=None, **kwargs):
#   def __init__(self, second_split=None, **kwargs):
#   def __init__(self, **kwargs):
# 
# Path: group_testing/simulator.py
# class Simulator(object):
#   def __init__(self,
#                workdir: str = None,
#                wetlab=wet_lab.WetLab(),
#                policy=group_policy.Dorfman(),
#                sampler=smc.SmcSampler(),
#                cheap_sampler=lbp.LbpSampler(),
#                num_simulations: int = 1,
#                num_tests_per_cycle: int = 10,
#                max_test_cycles: int = 5,
#                max_group_size: int = 8,
#                prior_specificity: Union[float, Iterable[float]] = 0.97,
#                prior_sensitivity: Union[float, Iterable[float]] = 0.85,
#                prior_infection_rate: Union[float, Iterable[float]] = 0.05,
#                metrics_cls=metrics.Metrics,
#                export_metrics_every: int = 5):
#   def reset(self, rng: np.ndarray):
#   def sampler(self):
#   def run(self, rngkey=None):
#   def is_over(self) -> bool:
#   def run_simulation(self, rng, sim_idx: int):
#   def run_one_cycle(self, rng, sim_idx: int):
#   def marginal(self):
#   def process_tests_results(self, rng, new_tests):
#   def get_new_groups(self, rng) -> np.ndarray:
#   def _resample(self, rng):
#   def _on_iteration_end(self,
#                         sim_idx: int,
#                         cycle_idx: int,
#                         new_groups=None,
#                         new_tests=None):
#   def interactive_loop(self, rngkey: int = None, show_fn=None):
# 
# Path: group_testing/wet_lab.py
# class WetLab:
# class WetLabCSV(WetLab):
#   def __init__(self,
#                num_patients: int = 2,
#                freeze_diseased: bool = True,
#                base_infection_rate: Union[float, Iterable[float]] = 0.05,
#                specificity: Union[float, Iterable[float]] = 0.97,
#                sensitivity: Union[float, Iterable[float]] = 0.85):
#   def reset(self, rng):
#   def group_tests_outputs(self, rng, groups):
#   def __init__(self, filename: str = None):
#   def reset(self, rng):
#   def group_tests_outputs(self, rng, groups):
#   def to_csv(self, groups: np.ndarray, marginals: np.ndarray):
#     def to_list(arr):
# 
# Path: group_testing/samplers/sequential_monte_carlo.py
# class SmcSampler(sampler.Sampler):
#   NAME = 'smc'
#   def __init__(self,
#                kernel=kernels.Gibbs(),
#                resample_at_each_iteration: bool = False,
#                start_from_prior: bool = True,
#                num_particles: int = 10000,
#                min_kernel_iterations: int = 2,
#                max_kernel_iterations: int = 20,
#                min_ratio_delta: float = 0.02,
#                target_unique_ratio: float = 0.95):
#   def reset(self):
#   def is_cheap(self):
#   def produce_sample(self, rng, state):
#   def resample_move(self,
#                     rng,
#                     particles,
#                     state,
#                     sampling_from_scratch):
#   def resample(self,
#                rng: int,
#                particle_weights: np.ndarray) -> np.ndarray:
#   def move(self,
#            rng: int,
#            particles: np.ndarray,
#            rho: float,
#            log_posterior_params: Dict[str, np.ndarray],
#            log_base_measure_params: Dict[str, np.ndarray]):


# Path: colab/interactive_loop.py
import time
import jax
import numpy as np
from group_testing import policy as group_policy
from group_testing import simulator
from group_testing import wet_lab
from group_testing.samplers import sequential_monte_carlo as smc

#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# Lint as: python3
"""An interactive loop to go back and forth from colab to spreadsheet."""





class InteractiveSimulator(simulator.Simulator):
  """Interactive group testing loop, reading an write into a spreadsheet."""

  def __init__(self, sheet):
    wetlab = wet_lab.WetLab(num_patients=sheet.num_patients)
    sampler = smc.SmcSampler(num_particles=sheet.num_patients * 150,
                             resample_at_each_iteration=True)

    policy_name = sheet.params['policy']
    if policy_name == 'G-MIMAX':

      policy = group_policy.MimaxPolicy(forward_iterations=2,