text: import logging
  from typing import Callable, Dict, Tuple, Optional
  
  import numpy as np
  import pyabc
  import torch
  import sbibm
  
  from sbi.inference import MCABC
  from sbibm.tasks.task import Task
  
  
  class PyAbcSimulator:
      """Wrapper from sbibm task to pyABC. 
  
      pyABC defines its own priors and they are sampled without batch dimension. This
      wrapper defines a call method that takes a single parameter set from a pyABC prior
      and uses the sbibm task simulator to generate the corresponding data and to return
      it in pyABC format. 
      """
  
      def __init__(self, task):
          self.simulator = task.get_simulator()
          self.dim_parameters = task.dim_parameters
          self.name = task.name
  
      def __call__(self, pyabc_parameter) -> Dict:
          parameters = torch.tensor(
              [[pyabc_parameter[f"param{dim+1}"] for dim in range(self.dim_parameters)]],
              dtype=torch.float32,
          )
          data = self.simulator(parameters).numpy().squeeze()
          return dict(data=data)
  
      @property
      def __name__(self) -> str:
          return self.name
  
  
  def wrap_prior(task):
      """Returns a pyABC.Distribution prior given a prior defined on a sbibm task.
  
      Note: works only for a specific set of priors: Uniform, LogNormal, Normal.
      """
      log = logging.getLogger(__name__)
      log.warn("Will discard any correlations in prior")
  
      bounds = {}
  
      prior_cls = str(task.prior_dist)
      if prior_cls == "Independent()":
          prior_cls = str(task.prior_dist.base_dist)
  
      prior_params = {}
      if "MultivariateNormal" in prior_cls:
          prior_params["m"] = task.prior_params["loc"].numpy()
          if "precision_matrix" in prior_cls:
              prior_params["C"] = np.linalg.inv(
                  task.prior_params["precision_matrix"].numpy()
              )
          if "covariance_matrix" in prior_cls:
              prior_params["C"] = task.prior_params["covariance_matrix"].numpy()
  
          prior_dict = {}
          for dim in range(task.dim_parameters):
              loc = prior_params["m"][dim]
              scale = np.sqrt(prior_params["C"][dim, dim])
  
              prior_dict[f"param{dim+1}"] = pyabc.RV("norm", loc, scale)
          prior = pyabc.Distribution(**prior_dict)
  
      elif "LogNormal" in prior_cls:
          # Note the difference in parameterisation between pytorch LogNormal and scipy
          # lognorm:
          # https://docs.scipy.org/doc/scipy/reference/generated/scipy.stats.lognorm.html#scipy.stats.lognorm
          prior_params["s"] = task.prior_params["scale"].numpy()
          prior_params["scale"] = np.exp(task.prior_params["loc"].numpy())
  
          prior_dict = {}
          for dim in range(task.dim_parameters):
              prior_dict[f"param{dim+1}"] = pyabc.RV(
                  "lognorm", s=prior_params["s"][dim], scale=prior_params["scale"][dim]
              )
  
          prior = pyabc.Distribution(**prior_dict)
  
      elif "Uniform" in prior_cls:
          prior_params["low"] = task.prior_params["low"].numpy()
          prior_params["high"] = task.prior_params["high"].numpy()
  
          prior_dict = {}
          for dim in range(task.dim_parameters):
              loc = prior_params["low"][dim]
              scale = prior_params["high"][dim] - loc
  
              prior_dict[f"param{dim+1}"] = pyabc.RV("uniform", loc, scale)
  
          prior = pyabc.Distribution(**prior_dict)
  
      else:
          log.info("No support for prior yet")
          raise NotImplementedError
  
      return prior
  
  
  def get_distance(distance: str) -> Callable:
      """Return distance function for pyabc."""
  
      if distance == "l1":
  
          def distance_fun(x, y):
              abs_diff = abs(x["data"] - y["data"])
              return np.atleast_1d(abs_diff).mean(axis=-1)
  
      elif distance == "mse":
  
          def distance_fun(x, y):
              return np.mean((x["data"] - y["data"]) ** 2, axis=-1)
  
      elif distance == "l2":
  
          def distance_fun(x, y):
              return np.linalg.norm(x["data"] - y["data"], axis=-1)
  
      else:
          raise NotImplementedError(f"Distance '{distance}' not implemented.")
  
      return distance_fun
  
  
  def clip_int(value, minimum, maximum):
      value = int(value)
      minimum = int(minimum)
      maximum = int(maximum)
      if value < minimum:
          return minimum
      if value > maximum:
          return maximum
      return value
  
  
  def run_pyabc(
      task: Task,
      db,
      num_simulations: int,
      observation: np.ndarray,
      pyabc_kwargs: dict,
      distance_str: str = "l2",
      batch_size: int = 1000,
      use_last_pop_samples: bool = False,
  ) -> Tuple[torch.Tensor, torch.Tensor]:
      """Run pyabc SMC with fixed budget and return particles and weights.
      
      Return previous population or prior samples if budget is exceeded.
      """
      log = sbibm.get_logger(__name__)
  
      abc = pyabc.ABCSMC(**pyabc_kwargs)
      abc.new(db, {"data": observation})
      history = abc.run(max_total_nr_simulations=num_simulations)
      num_calls = history.total_nr_simulations
  
      if num_calls < 1.0 * num_simulations:
          (particles_df, weights) = history.get_distribution(t=history.max_t)
          particles = torch.as_tensor(particles_df.values, dtype=torch.float32)
          weights = torch.as_tensor(weights, dtype=torch.float32)
      else:
          if history.max_t > 0:
              log.info(
                  f"Last population exceeded budget by {num_calls - num_simulations}."
              )
              (particles_df, weights) = history.get_distribution(t=history.max_t - 1)
              old_particles = torch.as_tensor(particles_df.values, dtype=torch.float32)
              old_weights = torch.as_tensor(weights, dtype=torch.float32)
              if use_last_pop_samples:
                  df = history.get_all_populations()
                  num_calls_last_pop = df.samples.values[-1]
                  over_budget = num_calls - num_simulations
                  proportion_over_budget = over_budget / num_calls_last_pop
                  # The proportion over budget needs to be replaced with old particles.
                  num_old_particles = int(
                      np.ceil(proportion_over_budget * pyabc_kwargs["population_size"])
                  )
                  log.info(
                      f"Filling up with {num_old_particles+1} samples from previous population."
                  )
                  # Combining populations.
                  (particles_df, weights) = history.get_distribution(t=history.max_t)
                  new_particles = torch.as_tensor(
                      particles_df.values, dtype=torch.float32
                  )
                  new_weights = torch.as_tensor(weights, dtype=torch.float32)
  
                  particles = torch.zeros_like(old_particles)
                  weights = torch.zeros_like(old_weights)
                  particles[:num_old_particles] = old_particles[:num_old_particles]
                  particles[num_old_particles:] = new_particles[num_old_particles:]
                  weights[:num_old_particles] = old_weights[:num_old_particles]
                  weights[num_old_particles:] = new_weights[num_old_particles:]
                  # Normalize combined weights.
                  weights /= weights.sum()
              else:
                  log.info("Returning previous population.")
                  particles = old_particles
                  weights = old_weights
          else:
              log.info("Running REJABC because first population exceeded budget.")
              posterior, _ = run_rejection_abc(
                  task,
                  num_simulations,
                  pyabc_kwargs["population_size"],
                  observation=torch.tensor(observation, dtype=torch.float32),
                  distance=distance_str,
                  batch_size=batch_size,
              )
              particles = posterior._samples
              weights = posterior._log_weights.exp()
  
      return particles, weights
  
  
  def run_rejection_abc(
      task: Task,
      num_simulations: int,
      population_size: int,
      observation: Optional[torch.Tensor] = None,
      distance: str = "l2",
      batch_size: int = 1000,
  ):
      """Return posterior and distances from a ABC with fixed budget."""
  
      inferer = MCABC(
          simulator=task.get_simulator(max_calls=num_simulations),
          prior=task.get_prior_dist(),
          simulation_batch_size=batch_size,
          distance=distance,
          show_progress_bars=True,
      )
      posterior, distances = inferer(
          x_o=observation,
          num_simulations=num_simulations,
          eps=None,
          quantile=population_size / num_simulations,
          return_distances=True,
      )
      return posterior, distances
  