class Classifier(abc.ABC):
    """An abstract base class for a ratio estimation classifier."""

    @abc.abstractmethod
    def __init__(self):
        """Initialize a classifier."""
        raise NotImplementedError

    @abc.abstractmethod
    def fit(self, X, y):
        """Fit a classifier.

        Parameters
        ----------
        X: np.ndarray (n_samples, n_features)
            Feature vectors of data.
        y: np.ndarray (n_samples, )
            Target values, must be binary.

        """
        raise NotImplementedError

    @abc.abstractmethod
    def predict_log_likelihood_ratio(self, X):
        """Predict a log-likelihood ratio.

        Parameters
        ----------
        X: np.ndarray (n_samples, n_features)
            Feature vectors of data.

        Returns
        -------
        np.ndarray

        """
        raise NotImplementedError

    def predict_likelihood_ratio(self, X):
        """Predict a likelihood ratio.

        Parameters
        ----------
        X: np.ndarray (n_samples, n_features)
            Feature vectors of data.

        Returns
        -------
        np.ndarray

        """
        return np.exp(self.predict_log_likelihood_ratio(X))

    @property
    @abc.abstractmethod
    def attributes(self):
        """Return attributes dictionary."""
        raise NotImplementedError class LogisticRegression(Classifier):
    """A logistic regression classifier for ratio estimation."""

    def __init__(self, config=None, class_min=0):
        """Initialize a logistic regression classifier."""
        self.config = self._resolve_config(config)
        self.class_min = self._resolve_class_min(class_min)
        self.model = LogReg(**self.config)
        self.scaler = StandardScaler()

    def fit(self, X, y):
        """Fit a logistic regression classifier."""
        Xs = self.scaler.fit_transform(X)
        self.model.fit(Xs, y)

    def predict_log_likelihood_ratio(self, X):
        """Predict a log-likelihood ratio."""
        Xs = self.scaler.transform(X)
        class_probs = np.maximum(self.model.predict_proba(Xs)[:, 1], self.class_min)
        return np.log(class_probs / (1 - class_probs))

    @property
    def attributes(self):
        """Return an attributes dictionary."""
        return {
            'parameters': {
                'coef_': self.model.coef_.tolist(),
                'intercept_': self.model.intercept_.tolist(),
                'n_iter': self.model.n_iter_.tolist()
            }
        }

    def _default_config(self):
        """Return a default config."""
        return {
            'penalty': 'l1',
            'solver': 'liblinear'
        }

    def _resolve_config(self, config):
        """Resolve a config for logistic regression classifier."""
        if not isinstance(config, dict):
            config = self._default_config()
        return config

    def _resolve_class_min(self, class_min):
        """Resolve a class min parameter that prevents negative inf values."""
        if isinstance(class_min, int) or isinstance(class_min, float):
            return class_min
        raise TypeError('class_min has to be either non-negative int or float') class LCBSC(AcquisitionBase):
    r"""Lower Confidence Bound Selection Criterion.

    Srinivas et al. call this GP-LCB.

    LCBSC uses the parameter delta which is here equivalent to 1/exploration_rate.

    Parameter delta should be in (0, 1) for the theoretical results to hold. The
    theoretical upper bound for total regret in Srinivas et al. has a probability greater
    or equal to 1 - delta, so values of delta very close to 1 or over it do not make much
    sense in that respect.

    Delta is roughly the exploitation tendency of the acquisition function.

    References
    ----------
    N. Srinivas, A. Krause, S. M. Kakade, and M. Seeger. Gaussian
    process optimization in the bandit setting: No regret and experimental design. In
    Proc. International Conference on Machine Learning (ICML), 2010

    E. Brochu, V.M. Cora, and N. de Freitas. A tutorial on Bayesian optimization of expensive
    cost functions, with application to active user modeling and hierarchical reinforcement
    learning. arXiv:1012.2599, 2010.

    Notes
    -----
    The formula presented in Brochu (pp. 15) seems to be from Srinivas et al. Theorem 2.
    However, instead of having t**(d/2 + 2) in \beta_t, it seems that the correct form
    would be t**(2d + 2).

    """

    def __init__(self, *args, delta=None, additive_cost=None, **kwargs):
        """Initialize LCBSC.

        Parameters
        ----------
        delta: float, optional
            In between (0, 1). Default is 1/exploration_rate. If given, overrides the
            exploration_rate.
        additive_cost: CostFunction, optional
            Cost function output is added to the base acquisition value.

        """
        if delta is not None:
            if delta <= 0 or delta >= 1:
                logger.warning('Parameter delta should be in the interval (0,1)')
            kwargs['exploration_rate'] = 1 / delta

        super(LCBSC, self).__init__(*args, **kwargs)
        self.name = 'lcbsc'
        self.label_fn = 'Confidence Bound'

        if additive_cost is not None and not isinstance(additive_cost, CostFunction):
            raise TypeError("Additive cost must be type CostFunction.")
        self.additive_cost = additive_cost

    @property
    def delta(self):
        """Return the inverse of exploration rate."""
        return 1 / self.exploration_rate

    def _beta(self, t):
        # Start from 0
        t += 1
        d = self.model.input_dim
        return 2 * np.log(t**(2 * d + 2) * np.pi**2 / (3 * self.delta))

    def evaluate(self, x, t=None):
        """Evaluate the Lower confidence bound selection criterion.

        Parameters
        ----------
        x: np.ndarray
        t: int, optional
            Current iteration (starting from 0).

        Returns
        -------
        np.ndarray

        """
        mean, var = self.model.predict(x, noiseless=True)
        value = mean - np.sqrt(self._beta(t) * var)
        if self.additive_cost is not None:
            value += self.additive_cost.evaluate(x)
        return value

    def evaluate_gradient(self, x, t=None):
        """Evaluate the gradient of the lower confidence bound selection criterion.

        Parameters
        ----------
        x: np.ndarray
        t: int, optional
            Current iteration (starting from 0).

        Returns
        -------
        np.ndarray

        """
        mean, var = self.model.predict(x, noiseless=True)
        grad_mean, grad_var = self.model.predictive_gradients(x)
        value = grad_mean - 0.5 * grad_var * np.sqrt(self._beta(t) / var)
        if self.additive_cost is not None:
            value += self.additive_cost.evaluate_gradient(x)
        return value class AcquisitionBase:
    """All acquisition functions are assumed to fulfill this interface.

    Gaussian noise ~N(0, self.noise_var) is added to the acquired points. By default,
    noise_var=0. You can define a different variance for the separate dimensions.

    """

    def __init__(self,
                 model,
                 prior=None,
                 n_inits=10,
                 max_opt_iters=1000,
                 noise_var=None,
                 exploration_rate=10,
                 seed=None,
                 constraints=None):
        """Initialize AcquisitionBase.

        Parameters
        ----------
        model : an object with attributes
                    input_dim : int
                    bounds : tuple of length 'input_dim' of tuples (min, max)
                and methods
                    evaluate(x) : function that returns model (mean, var, std)
        prior : scipy-like distribution, optional
            By default uniform distribution within model bounds.
        n_inits : int, optional
            Number of initialization points in internal optimization.
        max_opt_iters : int, optional
            Max iterations to optimize when finding the next point.
        noise_var : float or np.array, optional
            Acquisition noise variance for adding noise to the points near the optimized
            location. If array, must be 1d specifying the variance for different dimensions.
            Default: no added noise.
        exploration_rate : float, optional
            Exploration rate of the acquisition function (if supported)
        seed : int, optional
            Seed for getting consistent acquisition results. Used in getting random
            starting locations in acquisition function optimization.
        constraints : {Constraint, dict} or List of {Constraint, dict}, optional
            Additional model constraints.

        """
        self.model = model
        self.prior = prior
        self.n_inits = int(n_inits)
        self.max_opt_iters = int(max_opt_iters)
        self.constraints = constraints
        if noise_var is not None:
            self._check_noise_var(noise_var)
            self.noise_var = self._transform_noise_var(noise_var)
        else:
            self.noise_var = noise_var
        self.exploration_rate = exploration_rate
        self.random_state = np.random if seed is None else np.random.RandomState(seed)
        self.seed = 0 if seed is None else seed

    def _check_noise_var(self, noise_var):
        if isinstance(noise_var, dict):
            if not set(noise_var) == set(self.model.parameter_names):
                raise ValueError("Acquisition noise dictionary should contain all parameters.")

            if not all(isinstance(x, (int, float)) for x in noise_var.values()):
                raise ValueError("Acquisition noise dictionary values "
                                 "should all be int or float.")

            if any([x < 0 for x in noise_var.values()]):
                raise ValueError("Acquisition noises values should all be "
                                 "non-negative int or float.")

        elif isinstance(noise_var, (int, float)):
            if noise_var < 0:
                raise ValueError("Acquisition noise should be non-negative int or float.")
        else:
            raise ValueError("Either acquisition noise is a float or "
                             "it is a dictionary of floats defining "
                             "variance for each parameter dimension.")

    def _transform_noise_var(self, noise_var):
        if isinstance(noise_var, (float, int)):
            return noise_var

        # return a sorted list of noise variances in the same order than
        # parameter_names of the model
        if isinstance(noise_var, dict):
            return list(map(noise_var.get, self.model.parameter_names))

    def evaluate(self, x, t=None):
        """Evaluate the acquisition function at 'x'.

        Parameters
        ----------
        x : numpy.array
        t : int
            current iteration (starting from 0)

        """
        raise NotImplementedError

    def evaluate_gradient(self, x, t=None):
        """Evaluate the gradient of acquisition function at 'x'.

        Parameters
        ----------
        x : numpy.array
        t : int
            Current iteration (starting from 0).

        """
        raise NotImplementedError

    def acquire(self, n, t=None):
        """Return the next batch of acquisition points.

        Gaussian noise ~N(0, self.noise_var) is added to the acquired points.

        Parameters
        ----------
        n : int
            Number of acquisition points to return.
        t : int
            Current acq_batch_index (starting from 0).

        Returns
        -------
        x : np.ndarray
            The shape is (n, input_dim)

        """
        logger.debug('Acquiring the next batch of %d values', n)

        # Optimize the current minimum
        def obj(x):
            return self.evaluate(x, t)

        def grad_obj(x):
            return self.evaluate_gradient(x, t)

        xhat, _ = minimize(
            obj,
            self.model.bounds,
            method='L-BFGS-B' if self.constraints is None else 'SLSQP',
            constraints=self.constraints,
            grad=grad_obj,
            prior=self.prior,
            n_start_points=self.n_inits,
            maxiter=self.max_opt_iters,
            random_state=self.random_state)

        # Create n copies of the minimum
        x = np.tile(xhat, (n, 1))
        # Add noise for more efficient fitting of GP
        x = self._add_noise(x)

        return x

    def _add_noise(self, x):
        # Add noise for more efficient fitting of GP
        if self.noise_var is not None:
            noise_var = np.asanyarray(self.noise_var)
            if noise_var.ndim == 0:
                noise_var = np.tile(noise_var, self.model.input_dim)

            for i in range(self.model.input_dim):
                std = np.sqrt(noise_var[i])
                if std == 0:
                    continue
                xi = x[:, i]
                a = (self.model.bounds[i][0] - xi) / std
                b = (self.model.bounds[i][1] - xi) / std
                x[:, i] = ss.truncnorm.rvs(
                    a, b, loc=xi, scale=std, size=len(x), random_state=self.random_state)

        return x class GPyRegression:
    """Gaussian Process regression using the GPy library.

    GPy API: https://sheffieldml.github.io/GPy/
    """

    def __init__(self,
                 parameter_names=None,
                 bounds=None,
                 optimizer="scg",
                 max_opt_iters=50,
                 gp=None,
                 **gp_params):
        """Initialize GPyRegression.

        Parameters
        ----------
        parameter_names : list of str, optional
            Names of parameter nodes. If None, sets dimension to 1.
        bounds : dict, optional
            The region where to estimate the posterior for each parameter in
            model.parameters.
            `{'parameter_name':(lower, upper), ... }`
            If not supplied, defaults to (0, 1) bounds for all dimensions.
        optimizer : string, optional
            Optimizer for the GP hyper parameters
            Alternatives: "scg", "fmin_tnc", "simplex", "lbfgsb", "lbfgs", "sgd"
            See also: paramz.Model.optimize()
        max_opt_iters : int, optional
        gp : GPy.model.GPRegression instance, optional
        **gp_params
            kernel : GPy.Kern
            noise_var : float
            mean_function

        """
        if parameter_names is None:
            input_dim = 1
        elif isinstance(parameter_names, (list, tuple)):
            input_dim = len(parameter_names)
        else:
            raise ValueError("Keyword `parameter_names` must be a list of strings")

        if bounds is None:
            logger.warning('Parameter bounds not specified. Using [0,1] for each parameter.')
            bounds = [(0, 1)] * input_dim
        elif len(bounds) != input_dim:
            raise ValueError(
                'Length of `bounds` ({}) does not match the length of `parameter_names` ({}).'
                .format(len(bounds), input_dim))

        elif isinstance(bounds, dict):
            if len(bounds) == 1:  # might be the case parameter_names=None
                bounds = [bounds[n] for n in bounds.keys()]
            else:
                # turn bounds dict into a list in the same order as parameter_names
                bounds = [bounds[n] for n in parameter_names]
        else:
            raise ValueError("Keyword `bounds` must be a dictionary "
                             "`{'parameter_name': (lower, upper), ... }`")

        self.parameter_names = parameter_names
        self.input_dim = input_dim
        self.bounds = bounds

        self.gp_params = gp_params

        self.optimizer = optimizer
        self.max_opt_iters = max_opt_iters

        self._gp = gp

        self._rbf_is_cached = False
        self.is_sampling = False  # set to True once in sampling phase

    def __str__(self):
        """Return GPy's __str__."""
        return self._gp.__str__()

    def __repr__(self):
        """Return GPy's __str__."""
        return self.__str__()

    def predict(self, x, noiseless=False):
        """Return the GP model mean and variance at x.

        Parameters
        ----------
        x : np.array
            numpy compatible (n, input_dim) array of points to evaluate
            if len(x.shape) == 1 will be cast to 2D with x[None, :]
        noiseless : bool
            whether to include the noise variance or not to the returned variance

        Returns
        -------
        tuple
            GP (mean, var) at x where
                mean : np.array
                    with shape (x.shape[0], 1)
                var : np.array
                    with shape (x.shape[0], 1)

        """
        # Ensure it's 2d for GPy
        x = np.asanyarray(x).reshape((-1, self.input_dim))

        if self._gp is None:
            # TODO: return from GP mean function if given
            return np.zeros((x.shape[0], 1)), \
                np.ones((x.shape[0], 1))

        # direct (=faster) implementation for RBF kernel
        if self.is_sampling and self._kernel_is_default:
            if not self._rbf_is_cached:
                self._cache_RBF_kernel()

            r2 = np.sum(x**2., 1)[:, None] + self._rbf_x2sum - 2. * x.dot(self._gp.X.T)
            kx = self._rbf_var * np.exp(r2 * self._rbf_factor) + self._rbf_bias
            mu = kx.dot(self._rbf_woodbury)

            var = self._rbf_var + self._rbf_bias
            var -= kx.dot(self._rbf_woodbury_inv.dot(kx.T))
            var += self._rbf_noisevar  # likelihood

            return mu, var
        else:
            self._rbf_is_cached = False  # in case one resumes fitting the GP after sampling

        if noiseless:
            return self._gp.predict_noiseless(x)
        else:
            return self._gp.predict(x)

    # TODO: find a more general solution
    # cache some RBF-kernel-specific values for faster sampling
    def _cache_RBF_kernel(self):
        self._rbf_var = float(self._gp.kern.rbf.variance)
        self._rbf_factor = -0.5 / float(self._gp.kern.rbf.lengthscale)**2
        self._rbf_bias = float(self._gp.kern.bias.K(self._gp.X)[0, 0])
        self._rbf_noisevar = float(self._gp.likelihood.variance[0])
        self._rbf_woodbury = self._gp.posterior.woodbury_vector
        self._rbf_woodbury_inv = self._gp.posterior.woodbury_inv
        self._rbf_woodbury_chol = self._gp.posterior.woodbury_chol
        self._rbf_x2sum = np.sum(self._gp.X**2., 1)[None, :]
        self._rbf_is_cached = True

    def predict_mean(self, x):
        """Return the GP model mean function at x.

        Parameters
        ----------
        x : np.array
            numpy compatible (n, input_dim) array of points to evaluate
            if len(x.shape) == 1 will be cast to 2D with x[None, :]

        Returns
        -------
        np.array
            with shape (x.shape[0], 1)

        """
        return self.predict(x)[0]

    def predictive_gradients(self, x):
        """Return the gradients of the GP model mean and variance at x.

        Parameters
        ----------
        x : np.array
            numpy compatible (n, input_dim) array of points to evaluate
            if len(x.shape) == 1 will be cast to 2D with x[None, :]

        Returns
        -------
        tuple
            GP (grad_mean, grad_var) at x where
                grad_mean : np.array
                    with shape (x.shape[0], input_dim)
                grad_var : np.array
                    with shape (x.shape[0], input_dim)

        """
        # Ensure it's 2d for GPy
        x = np.asanyarray(x).reshape((-1, self.input_dim))

        if self._gp is None:
            # TODO: return from GP mean function if given
            return np.zeros((x.shape[0], self.input_dim)), \
                np.zeros((x.shape[0], self.input_dim))

        # direct (=faster) implementation for RBF kernel
        if self.is_sampling and self._kernel_is_default:
            if not self._rbf_is_cached:
                self._cache_RBF_kernel()

            r2 = np.sum(x**2., 1)[:, None] + self._rbf_x2sum - 2. * x.dot(self._gp.X.T)
            kx = self._rbf_var * np.exp(r2 * self._rbf_factor)
            dkdx = 2. * self._rbf_factor * (x - self._gp.X) * kx.T
            grad_mu = dkdx.T.dot(self._rbf_woodbury).T

            v = np.linalg.solve(self._rbf_woodbury_chol, kx.T + self._rbf_bias)
            dvdx = np.linalg.solve(self._rbf_woodbury_chol, dkdx)
            grad_var = -2. * dvdx.T.dot(v).T
        else:
            grad_mu, grad_var = self._gp.predictive_gradients(x)
            grad_mu = grad_mu[:, :, 0]  # Assume 1D output (distance in ABC)

        return grad_mu, grad_var

    def predictive_gradient_mean(self, x):
        """Return the gradient of the GP model mean at x.

        Parameters
        ----------
        x : np.array
            numpy compatible (n, input_dim) array of points to evaluate
            if len(x.shape) == 1 will be cast to 2D with x[None, :]

        Returns
        -------
        np.array
            with shape (x.shape[0], input_dim)

        """
        return self.predictive_gradients(x)[0]

    def _init_gp(self, x, y):
        self._kernel_is_default = False

        if self.gp_params.get('kernel') is None:
            kernel = self._default_kernel(x, y)

            if self.gp_params.get('noise_var') is None and self.gp_params.get(
                    'mean_function') is None:
                self._kernel_is_default = True

        else:
            kernel = self.gp_params.get('kernel')

        noise_var = self.gp_params.get('noise_var') or np.max(y)**2. / 100.
        mean_function = self.gp_params.get('mean_function')
        self._gp = self._make_gpy_instance(
            x, y, kernel=kernel, noise_var=noise_var, mean_function=mean_function)

    def _default_kernel(self, x, y):
        # Some heuristics to choose kernel parameters based on the initial data
        length_scale = (np.max(self.bounds) - np.min(self.bounds)) / 3.
        kernel_var = (np.max(y) / 3.)**2.
        bias_var = kernel_var / 4.

        # Construct a default kernel
        kernel = GPy.kern.RBF(input_dim=self.input_dim)

        # Set the priors
        kernel.lengthscale.set_prior(
            GPy.priors.Gamma.from_EV(length_scale, length_scale), warning=False)
        kernel.variance.set_prior(GPy.priors.Gamma.from_EV(kernel_var, kernel_var), warning=False)

        # If no mean function is specified, add a bias term to the kernel
        if 'mean_function' not in self.gp_params:
            bias = GPy.kern.Bias(input_dim=self.input_dim)
            bias.set_prior(GPy.priors.Gamma.from_EV(bias_var, bias_var), warning=False)
            kernel += bias

        return kernel

    def _make_gpy_instance(self, x, y, kernel, noise_var, mean_function):
        return GPy.models.GPRegression(
            X=x, Y=y, kernel=kernel, noise_var=noise_var, mean_function=mean_function)

    def update(self, x, y, optimize=False):
        """Update the GP model with new data.

        Parameters
        ----------
        x : np.array
        y : np.array
        optimize : bool, optional
            Whether to optimize hyperparameters.

        """
        # Must cast these as 2d for GPy
        x = x.reshape((-1, self.input_dim))
        y = y.reshape((-1, 1))

        if self._gp is None:
            self._init_gp(x, y)
        else:
            # Reconstruct with new data
            x = np.r_[self._gp.X, x]
            y = np.r_[self._gp.Y, y]
            # It seems that GPy will do some optimization unless you make copies of everything
            kernel = self._gp.kern.copy() if self._gp.kern else None
            noise_var = self._gp.Gaussian_noise.variance[0]
            mean_function = self._gp.mean_function.copy() if self._gp.mean_function else None
            self._gp = self._make_gpy_instance(
                x, y, kernel=kernel, noise_var=noise_var, mean_function=mean_function)

        if optimize:
            self.optimize()

    def optimize(self):
        """Optimize GP hyperparameters."""
        logger.debug("Optimizing GP hyperparameters")
        try:
            self._gp.optimize(self.optimizer, max_iters=self.max_opt_iters)
        except np.linalg.linalg.LinAlgError:
            logger.warning("Numerical error in GP optimization. Stopping optimization")

    @property
    def n_evidence(self):
        """Return the number of observed samples."""
        if self._gp is None:
            return 0
        return self._gp.num_data

    @property
    def X(self):
        """Return input evidence."""
        return self._gp.X

    @property
    def Y(self):
        """Return output evidence."""
        return self._gp.Y

    @property
    def noise(self):
        """Return the noise."""
        return self._gp.Gaussian_noise.variance[0]

    @property
    def instance(self):
        """Return the gp instance."""
        return self._gp

    def copy(self):
        """Return a copy of current instance."""
        kopy = copy.copy(self)
        if self._gp:
            kopy._gp = self._gp.copy()

        if 'kernel' in self.gp_params:
            kopy.gp_params['kernel'] = self.gp_params['kernel'].copy()

        if 'mean_function' in self.gp_params:
            kopy.gp_params['mean_function'] = self.gp_params['mean_function'].copy()

        return kopy

    def __copy__(self):
        """Return a copy of current instance."""
        return self.copy() class CostFunction:
    """Convenience class for modelling acquisition costs."""

    def __init__(self, function, gradient, scale=1):
        """Initialise CostFunction.

        Parameters
        ----------
        function : callable
            Function that returns cost function value.
        gradient : callable
            Function that returns cost function gradient.
        scale : float, optional
            Cost function is multiplied with scale.

        """
        self.function = function
        self.gradient = gradient
        self.scale = scale

    def evaluate(self, x):
        """Return cost function value evaluated at x.

        Parameters
        ----------
        x : np.ndarray, shape: (input_dim,) or (n, input_dim)

        Returns
        -------
        np.ndarray, shape: (n, 1)

        """
        x = np.atleast_2d(x)
        n, input_dim = x.shape
        return self.scale * self.function(x).reshape(n, 1)

    def evaluate_gradient(self, x):
        """Return cost function gradient evaluated at x.

        Parameters
        ----------
        x : np.ndarray, shape: (input_dim,) or (n, input_dim)

        Returns
        -------
        np.ndarray, shape: (n, input_dim)

        """
        x = np.atleast_2d(x)
        n, input_dim = x.shape
        return self.scale * self.gradient(x).reshape(n, input_dim) class ParameterInference:
    """A base class for parameter inference methods.

    Attributes
    ----------
    model : elfi.ElfiModel
        The ELFI graph used by the algorithm
    output_names : list
        Names of the nodes whose outputs are included in the batches
    client : elfi.client.ClientBase
        The batches are computed in the client
    max_parallel_batches : int
    state : dict
        Stores any changing data related to achieving the objective. Must include a key
        ``n_batches`` for determining when the inference is finished.
    objective : dict
        Holds the data for the algorithm to internally determine how many batches are still
        needed. You must have a key ``n_batches`` here. By default the algorithm finished when
        the ``n_batches`` in the state dictionary is equal or greater to the corresponding
        objective value.
    batches : elfi.client.BatchHandler
        Helper class for submitting batches to the client and keeping track of their
        indexes.
    pool : elfi.store.OutputPool
        Pool object for storing and reusing node outputs.


    """

    def __init__(self,
                 model,
                 output_names,
                 batch_size=1,
                 seed=None,
                 pool=None,
                 max_parallel_batches=None):
        """Construct the inference algorithm object.

        If you are implementing your own algorithm do not forget to call `super`.

        Parameters
        ----------
        model : ElfiModel
            Model to perform the inference with.
        output_names : list
            Names of the nodes whose outputs will be requested from the ELFI graph.
        batch_size : int, optional
            The number of parameter evaluations in each pass through the ELFI graph.
            When using a vectorized simulator, using a suitably large batch_size can provide
            a significant performance boost.
        seed : int, optional
            Seed for the data generation from the ElfiModel
        pool : OutputPool, optional
            OutputPool both stores and provides precomputed values for batches.
        max_parallel_batches : int, optional
            Maximum number of batches allowed to be in computation at the same time.
            Defaults to number of cores in the client


        """
        model = model.model if isinstance(model, NodeReference) else model
        if not model.parameter_names:
            raise ValueError('Model {} defines no parameters'.format(model))

        self.model = model.copy()
        self.output_names = self._check_outputs(output_names)

        self.client = elfi.client.get_client()

        # Prepare the computation_context
        context = ComputationContext(
            batch_size=batch_size, seed=seed, pool=pool)
        self.batches = elfi.client.BatchHandler(
            self.model, context=context, output_names=output_names, client=self.client)
        self.computation_context = context
        self.max_parallel_batches = max_parallel_batches or self.client.num_cores

        if self.max_parallel_batches <= 0:
            msg = 'Value for max_parallel_batches ({}) must be at least one.'.format(
                self.max_parallel_batches)
            if self.client.num_cores == 0:
                msg += ' Client has currently no workers available. Please make sure ' \
                       'the cluster has fully started or set the max_parallel_batches ' \
                       'parameter by hand.'
            raise ValueError(msg)

        # State and objective should contain all information needed to continue the
        # inference after an iteration.
        self.state = dict(n_sim=0, n_batches=0)
        self.objective = dict()
        self.progress_bar = ProgressBar(prefix='Progress', suffix='Complete',
                                        decimals=1, length=50, fill='=')

    @property
    def pool(self):
        """Return the output pool of the inference."""
        return self.computation_context.pool

    @property
    def seed(self):
        """Return the seed of the inference."""
        return self.computation_context.seed

    @property
    def parameter_names(self):
        """Return the parameters to be inferred."""
        return self.model.parameter_names

    @property
    def batch_size(self):
        """Return the current batch_size."""
        return self.computation_context.batch_size

    def set_objective(self, *args, **kwargs):
        """Set the objective of the inference.

        This method sets the objective of the inference (values typically stored in the
        `self.objective` dict).

        Returns
        -------
        None

        """
        raise NotImplementedError

    def extract_result(self):
        """Prepare the result from the current state of the inference.

        ELFI calls this method in the end of the inference to return the result.

        Returns
        -------
        result : elfi.methods.result.Result

        """
        raise NotImplementedError

    def update(self, batch, batch_index):
        """Update the inference state with a new batch.

        ELFI calls this method when a new batch has been computed and the state of
        the inference should be updated with it. It is also possible to bypass ELFI and
        call this directly to update the inference.

        Parameters
        ----------
        batch : dict
            dict with `self.outputs` as keys and the corresponding outputs for the batch
            as values
        batch_index : int

        Returns
        -------
        None

        """
        self.state['n_batches'] += 1
        self.state['n_sim'] += self.batch_size

    def prepare_new_batch(self, batch_index):
        """Prepare values for a new batch.

        ELFI calls this method before submitting a new batch with an increasing index
        `batch_index`. This is an optional method to override. Use this if you have a need
        do do preparations, e.g. in Bayesian optimization algorithm, the next acquisition
        points would be acquired here.

        If you need provide values for certain nodes, you can do so by constructing a
        batch dictionary and returning it. See e.g. BayesianOptimization for an example.

        Parameters
        ----------
        batch_index : int
            next batch_index to be submitted

        Returns
        -------
        batch : dict or None
            Keys should match to node names in the model. These values will override any
            default values or operations in those nodes.

        """
        pass

    def plot_state(self, **kwargs):
        """Plot the current state of the algorithm.

        Parameters
        ----------
        axes : matplotlib.axes.Axes (optional)
        figure : matplotlib.figure.Figure (optional)
        xlim
            x-axis limits
        ylim
            y-axis limits
        interactive : bool (default False)
            If true, uses IPython.display to update the cell figure
        close
            Close figure in the end of plotting. Used in the end of interactive mode.

        Returns
        -------
        None

        """
        raise NotImplementedError

    def infer(self, *args, vis=None, bar=True, **kwargs):
        """Set the objective and start the iterate loop until the inference is finished.

        See the other arguments from the `set_objective` method.

        Parameters
        ----------
        vis : dict, optional
            Plotting options. More info in self.plot_state method
        bar : bool, optional
            Flag to remove (False) or keep (True) the progress bar from/in output.

        Returns
        -------
        result : Sample

        """
        vis_opt = vis if isinstance(vis, dict) else {}

        self.set_objective(*args, **kwargs)

        while not self.finished:
            self.iterate()
            if vis:
                self.plot_state(interactive=True, **vis_opt)

            if bar:
                self.progress_bar.update_progressbar(self.state['n_batches'],
                                                     self._objective_n_batches)

        self.batches.cancel_pending()
        if vis:
            self.plot_state(close=True, **vis_opt)

        return self.extract_result()

    def iterate(self):
        """Advance the inference by one iteration.

        This is a way to manually progress the inference. One iteration consists of
        waiting and processing the result of the next batch in succession and possibly
        submitting new batches.

        Notes
        -----
        If the next batch is ready, it will be processed immediately and no new batches
        are submitted.

        New batches are submitted only while waiting for the next one to complete. There
        will never be more batches submitted in parallel than the `max_parallel_batches`
        setting allows.

        Returns
        -------
        None

        """
        # Submit new batches if allowed
        while self._allow_submit(self.batches.next_index):
            next_batch = self.prepare_new_batch(self.batches.next_index)
            logger.debug("Submitting batch %d" % self.batches.next_index)
            self.batches.submit(next_batch)

        # Handle the next ready batch in succession
        batch, batch_index = self.batches.wait_next()
        logger.debug('Received batch %d' % batch_index)
        self.update(batch, batch_index)

    @property
    def finished(self):
        """Check whether objective of n_batches have been reached."""
        return self._objective_n_batches <= self.state['n_batches']

    def _allow_submit(self, batch_index):
        return (self.max_parallel_batches > self.batches.num_pending
                and self._has_batches_to_submit and (not self.batches.has_ready()))

    @property
    def _has_batches_to_submit(self):
        return self._objective_n_batches > self.state['n_batches'] + self.batches.num_pending

    @property
    def _objective_n_batches(self):
        """Check that n_batches can be computed from the objective."""
        if 'n_batches' in self.objective:
            n_batches = self.objective['n_batches']
        elif 'n_sim' in self.objective:
            n_batches = ceil(self.objective['n_sim'] / self.batch_size)
        else:
            raise ValueError(
                'Objective must define either `n_batches` or `n_sim`.')
        return n_batches

    def _extract_result_kwargs(self):
        """Extract common arguments for the ParameterInferenceResult object."""
        return {
            'method_name': self.__class__.__name__,
            'parameter_names': self.parameter_names,
            'seed': self.seed,
            'n_sim': self.state['n_sim'],
            'n_batches': self.state['n_batches']
        }

    @staticmethod
    def _resolve_model(model, target, default_reference_class=NodeReference):
        if isinstance(model, ElfiModel) and target is None:
            raise NotImplementedError(
                "Please specify the target node of the inference method")

        if isinstance(model, NodeReference):
            target = model
            model = target.model

        if isinstance(target, str):
            target = model[target]

        if not isinstance(target, default_reference_class):
            raise ValueError('Unknown target node class')

        return model, target.name

    def _check_outputs(self, output_names):
        """Filter out duplicates and check that corresponding nodes exist.

        Preserves the order.
        """
        output_names = output_names or []
        checked_names = []
        seen = set()
        for name in output_names:
            if isinstance(name, NodeReference):
                name = name.name

            if name in seen:
                continue
            elif not isinstance(name, str):
                raise ValueError(
                    'All output names must be strings, object {} was given'.format(name))
            elif not self.model.has_node(name):
                raise ValueError(
                    'Node {} output was requested, but it is not in the model.')

            seen.add(name)
            checked_names.append(name)

        return checked_names class BOLFIREPosterior:
    """Results from BOLFIRE inference method."""

    def __init__(self,
                 parameter_names,
                 model,
                 prior,
                 classifier_attributes,
                 n_opt_inits=10,
                 max_opt_iters=1000,
                 *args, **kwargs):
        """Initialize BOLFIREPosterior.

        Parameters
        ----------
        parameter_names: list
            Names of the parameter nodes.
        model: elfi.bo.gpy_regression.GPyRegression
            Instance of the surrogate model
        prior: elfi.methods.utils.ModelPrior
            Joint prior of the elfi model.
        classifier_attributes: list
            Classifier's attributes on each inference round.
        n_opt_inits: int, optional
            Number of initialization points in internal optimization.
        max_opt_iters: int, optional
            Maximum number of iterations performed in internal optimization.

        """
        self._parameter_names = parameter_names
        self._model = model
        self._prior = prior
        self._n_opt_inits = n_opt_inits
        self._max_opt_iters = max_opt_iters
        self._classifier_attributes = classifier_attributes

        # compute map estimates
        self._map_estimates = self._compute_map_estimates()

    @property
    def map_estimates(self):
        """Return the maximum a posterior estimate for each parameter.

        Returns
        -------
        OrderedDict

        """
        return OrderedDict([(parameter_name, self._map_estimates[i]) for i, parameter_name
                            in enumerate(self._parameter_names)])

    @property
    def classifier_attributes(self):
        """Return the classifier's attributes."""
        return self._classifier_attributes

    @property
    def surrogate_model_attributes(self):
        """Return the surrogate model's (GP regression) attributes."""
        return {
            'parameters': self._model._gp.param_array.tolist(),
            'X': self._model.X.tolist(),
            'Y': self._model.Y.tolist()
        }

    def pdf(self, x):
        """Return the unnormalized posterior at x.

        Parameters
        ----------
        x: np.ndarray

        Returns
        -------
        np.ndarray

        """
        return np.exp(self.logpdf(x))

    def _negative_pdf(self, x):
        """Return the negative unnormalized posterior at x."""
        return -1 * self.pdf(x)

    def logpdf(self, x):
        """Return the unnormalized log-posterior at x.

        Parameters
        ----------
        x: np.ndarray

        Returns
        -------
        np.ndarray

        """
        return self._prior.logpdf(x).reshape(-1, 1) - self._model.predict_mean(x)

    def gradient_pdf(self, x):
        """Return the gradient of the unnormalized posterior pdf at x.

        Parameters
        ----------
        x: np.ndarray

        Returns
        -------
        np.ndarray

        """
        return np.exp(self.logpdf(x)) * self.gradient_logpdf(x)

    def _negative_gradient_pdf(self, x):
        """Return the negative gradient of the unnormalized posterior pdf at x."""
        return -1 * self.gradient_pdf(x)

    def gradient_logpdf(self, x):
        """Return the gradient of unnormalized log-posterior pdf at x.

        Parameters
        ----------
        x: np.ndarray

        Returns
        -------
        np.ndarray

        """
        return self._prior.gradient_logpdf(x).reshape(1, -1) \
            - self._model.predictive_gradient_mean(x)

    def _compute_map_estimates(self):
        """Return the maximum a posterior estimate for each parameter."""
        minimum_location, _ = minimize(
            fun=self._negative_pdf,
            bounds=self._model.bounds,
            grad=self._negative_gradient_pdf,
            prior=self._prior,
            n_start_points=self._n_opt_inits,
            maxiter=self._max_opt_iters
        )
        return minimum_location class BOLFIRESample(Sample):
    """Container for results from BOLFIRE."""

    def __init__(self, method_name, chains, parameter_names, warmup, *args, **kwargs):
        """Initialize BOLFIRE result.

        Parameters
        ----------
        method_name: str
            Name of the inference method.
        chains: np.ndarray (n_chains, n_samples, n_parameters)
            Chains from sampling, warmup included.
        parameter_names: list
            List of names in the outputs dict that refer to model parameters.
        warmup: int
            Number of warmup iterations in chains.

        """
        n_chains = chains.shape[0]
        warmed_up = chains[:, warmup:, :]
        concatenated = warmed_up.reshape((-1,) + chains.shape[2:])
        outputs = dict(zip(parameter_names, concatenated.T))

        super(BOLFIRESample, self).__init__(
            method_name=method_name,
            outputs=outputs,
            parameter_names=parameter_names,
            chains=chains,
            n_chains=n_chains,
            warmup=warmed_up,
            *args, **kwargs
        ) def arr2d_to_batch(x, names):
    """Convert a 2d array to a batch dictionary columnwise.

    Parameters
    ----------
    x : np.ndarray
        2d array of values
    names : list[str]
        List of names

    Returns
    -------
    dict
        A batch dictionary

    """
    # TODO: support vector parameter nodes
    try:
        x = x.reshape((-1, len(names)))
    except BaseException:
        raise ValueError("A dimension mismatch in converting array to batch dictionary. "
                         "This may be caused by multidimensional "
                         "prior nodes that are not yet supported.")
    batch = {p: x[:, i] for i, p in enumerate(names)}
    return batch def batch_to_arr2d(batches, names):
    """Convert batches into a single numpy array.

    Parameters
    ----------
    batches : dict or list
       A list of batches or a single batch
    names : list
       Name of outputs to include in the array. Specifies the order.

    Returns
    -------
    np.array
        2d, where columns are batch outputs

    """
    if not batches:
        return []
    if not isinstance(batches, list):
        batches = [batches]

    rows = []
    for batch_ in batches:
        rows.append(np.column_stack([batch_[n] for n in names]))

    return np.vstack(rows) def resolve_sigmas(parameter_names: List[str],
                   sigma_proposals: Optional[Dict] = None,
                   bounds: Optional[Dict] = None) -> List:
    """Map dictionary of sigma_proposals into a list order as parameter_names.

    Parameters
    ----------
    parameter_names: List[str]
        names of the parameters
    sigma_proposals: Dict
        non-negative standard deviations for each dimension
        {'parameter_name': float}
    bounds : Dict, optional
        the region where to estimate the posterior for each parameter in
        model.parameters
        `{'parameter_name':(lower, upper), ... }

    Returns
    -------
    List
       list of sigma_proposals in the same order than in parameter_names

    """
    if sigma_proposals is None:
        sigma_proposals = []
        for bound in bounds:
            length_interval = bound[1] - bound[0]
            sigma_proposals.append(length_interval / 10)
    elif isinstance(sigma_proposals, dict):
        errmsg = "sigma_proposals' keys have to be identical to " \
                    "target_model.parameter_names."
        if len(sigma_proposals) is not len(parameter_names):
            raise ValueError(errmsg)
        try:
            sigma_proposals = [sigma_proposals[x] for x in parameter_names]
        except ValueError:
            print(parameter_names)
    else:
        raise ValueError("If provided, sigma_proposals need to be input as a dict.")

    return sigma_proposals class ElfiModel(GraphicalModel):
    """A container for the inference model.

    The ElfiModel is a directed acyclic graph (DAG), whose nodes represent
    parts of the inference task, for example the parameters to be inferred,
    the simulator or a summary statistic.
    """

    def __init__(self, name=None, observed=None, source_net=None):
        """Initialize the inference model.

        Parameters
        ----------
        name : str, optional
        observed : dict, optional
            Observed data with node names as keys.
        source_net : nx.DiGraph, optional
        set_current : bool, optional
            Sets this model as the current (default) ELFI model

        """
        super(ElfiModel, self).__init__(source_net)
        self.name = name or "model_{}".format(random_name())
        self.observed = observed or {}

    @property
    def name(self):
        """Return name of the model."""
        return self.source_net.graph['name']

    @name.setter
    def name(self, name):
        """Set the name of the model."""
        self.source_net.graph['name'] = name

    @property
    def observed(self):
        """Return the observed data for the nodes in a dictionary."""
        return self.source_net.graph['observed']

    @observed.setter
    def observed(self, observed):
        """Set the observed data of the model.

        Parameters
        ----------
        observed : dict

        """
        if not isinstance(observed, dict):
            raise ValueError("Observed data must be given in a dictionary with the node"
                             "name as the key")
        self.source_net.graph['observed'] = observed

    def generate(self, batch_size=1, outputs=None, with_values=None, seed=None):
        """Generate a batch of outputs.

        This method is useful for testing that the ELFI graph works.

        Parameters
        ----------
        batch_size : int, optional
        outputs : list, optional
        with_values : dict, optional
            You can specify values for nodes to use when generating data
        seed : int, optional
            Defaults to global numpy seed.

        """
        if outputs is None:
            outputs = self.source_net.nodes()
        elif isinstance(outputs, str):
            outputs = [outputs]
        if not isinstance(outputs, list):
            raise ValueError('Outputs must be a list of node names')

        if seed is None:
            seed = 'global'

        pool = None
        if with_values is not None:
            pool = OutputPool(with_values.keys())
            pool.add_batch(with_values, 0)
        context = ComputationContext(batch_size, seed=seed, pool=pool)

        client = elfi.client.get_client()
        compiled_net = client.compile(self.source_net, outputs)
        loaded_net = client.load_data(compiled_net, context, batch_index=0)
        return client.compute(loaded_net)

    def get_reference(self, name):
        """Return a new reference object for a node in the model.

        Parameters
        ----------
        name : str

        """
        cls = self.get_node(name)['attr_dict']['_class']
        return cls.reference(name, self)

    def get_state(self, name):
        """Return the state of the node.

        Parameters
        ----------
        name : str

        """
        return self.source_net.nodes[name]

    def update_node(self, name, updating_name):
        """Update `node` with `updating_node` in the model.

        The node with name `name` gets the state (operation), parents and observed
        data (if applicable) of the updating_node. The updating node is then removed
        from the graph.

        Parameters
        ----------
        name : str
        updating_name : str

        """
        update_observed = False
        obs = None
        if updating_name in self.observed:
            update_observed = True
            obs = self.observed.pop(updating_name)

        super(ElfiModel, self).update_node(name, updating_name)

        # Move data to the updated node
        if update_observed:
            self.observed[name] = obs

    def remove_node(self, name):
        """Remove a node from the graph.

        Parameters
        ----------
        name : str

        """
        if name in self.observed:
            self.observed.pop(name)
        super(ElfiModel, self).remove_node(name)

    @property
    def parameter_names(self):
        """Return a list of model parameter names in an alphabetical order."""
        return sorted([n for n in self.nodes if '_parameter' in self.get_state(n)['attr_dict']])

    @parameter_names.setter
    def parameter_names(self, parameter_names):
        """Set the model parameter nodes.

        For each node name in parameters, the corresponding node will be marked as being a
        parameter node. Other nodes will be marked as not being parameter nodes.

        Parameters
        ----------
        parameter_names : list
            A list of parameter names

        """
        parameter_names = set(parameter_names)
        for n in self.nodes:
            state = self.get_state(n)['attr_dict']
            if n in parameter_names:
                parameter_names.remove(n)
                state['_parameter'] = True
            else:
                if '_parameter' in state:
                    state.pop('_parameter')
        if len(parameter_names) > 0:
            raise ValueError('Parameters {} not found from the model'.format(parameter_names))

    def copy(self):
        """Return a copy of the ElfiModel instance.

        Returns
        -------
        ElfiModel

        """
        kopy = super(ElfiModel, self).copy()
        kopy.name = "{}_copy_{}".format(self.name, random_name())
        return kopy

    def save(self, prefix=None):
        """Save the current model to pickled file.

        Parameters
        ----------
        prefix : str, optional
            Path to the directory under which to save the model. Default is the current working
            directory.

        """
        path = self.name + '.pkl'
        if prefix is not None:
            os.makedirs(prefix, exist_ok=True)
            path = os.path.join(prefix, path)
        pickle.dump(self, open(path, "wb"))

    @classmethod
    def load(cls, name, prefix):
        """Load the pickled ElfiModel.

        Assumes there exists a file "name.pkl" in the current directory.

        Parameters
        ----------
        name : str
            Name of the model file to load (without the .pkl extension).
        prefix : str
            Path to directory where the model file is located, optional.

        Returns
        -------
        ElfiModel

        """
        path = name + '.pkl'
        if prefix is not None:
            path = os.path.join(prefix, path)
        return pickle.load(open(path, "rb"))

    def __getitem__(self, node_name):
        """Return a new reference object for a node in the model.

        Parameters
        ----------
        node_name : str

        """
        return self.get_reference(node_name) class Summary(ObservableMixin, NodeReference):
    """A summary node of an ELFI graph.

    Summary nodes are deterministic operations associated with the observed data. if their
    parents hold observed data it will be automatically transformed.
    """

    def __init__(self, fn, *parents, **kwargs):
        """Initialize a Summary.

        Parameters
        ----------
        fn : callable
            Summary function with a signature `summary(*parents)`
        parents
            Input data for the summary function.
        kwargs

        """
        if not parents:
            raise ValueError('This node requires that at least one parent is specified.')
        state = dict(_operation=fn)
        super(Summary, self).__init__(*parents, state=state, **kwargs) class ModelPrior:
    """Construct a joint prior distribution over all the parameter nodes in `ElfiModel`."""

    def __init__(self, model):
        """Initialize a ModelPrior.

        Parameters
        ----------
        model : ElfiModel

        """
        model = model.copy()
        self.parameter_names = model.parameter_names
        self.dim = len(self.parameter_names)
        self.client = Client()

        # Prepare nets for the pdf methods
        self._pdf_node = augmenter.add_pdf_nodes(model, log=False)[0]
        self._logpdf_node = augmenter.add_pdf_nodes(model, log=True)[0]

        self._rvs_net = self.client.compile(model.source_net, outputs=self.parameter_names)
        self._pdf_net = self.client.compile(model.source_net, outputs=self._pdf_node)
        self._logpdf_net = self.client.compile(model.source_net, outputs=self._logpdf_node)

    def rvs(self, size=None, random_state=None):
        """Sample the joint prior."""
        random_state = np.random if random_state is None else random_state

        context = ComputationContext(size or 1, seed='global')
        loaded_net = self.client.load_data(self._rvs_net, context, batch_index=0)

        # Change to the correct random_state instance
        # TODO: allow passing random_state to ComputationContext seed
        loaded_net.nodes['_random_state'].update({'output': random_state})
        del loaded_net.nodes['_random_state']['operation']

        batch = self.client.compute(loaded_net)
        rvs = np.column_stack([batch[p] for p in self.parameter_names])

        if self.dim == 1:
            rvs = rvs.reshape(size or 1)

        return rvs[0] if size is None else rvs

    def pdf(self, x):
        """Return the density of the joint prior at x."""
        return self._evaluate_pdf(x)

    def logpdf(self, x):
        """Return the log density of the joint prior at x."""
        return self._evaluate_pdf(x, log=True)

    def _evaluate_pdf(self, x, log=False):
        if log:
            net = self._logpdf_net
            node = self._logpdf_node
        else:
            net = self._pdf_net
            node = self._pdf_node

        x = np.asanyarray(x)
        ndim = x.ndim
        x = x.reshape((-1, self.dim))
        batch = self._to_batch(x)

        # TODO: we could add a seed value that would load a "random state" instance
        #       throwing an error if it is used, for instance seed="not used".
        context = ComputationContext(len(x), seed=0)
        loaded_net = self.client.load_data(net, context, batch_index=0)

        # Override
        for k, v in batch.items():
            loaded_net.nodes[k].update({'output': v})
            del loaded_net.nodes[k]['operation']

        val = self.client.compute(loaded_net)[node]
        if ndim == 0 or (ndim == 1 and self.dim > 1):
            val = val[0]

        return val

    def gradient_pdf(self, x):
        """Return the gradient of density of the joint prior at x."""
        raise NotImplementedError

    def gradient_logpdf(self, x, stepsize=None):
        """Return the gradient of log density of the joint prior at x.

        Parameters
        ----------
        x : float or np.ndarray
        stepsize : float or list
            Stepsize or stepsizes for the dimensions

        """
        x = np.asanyarray(x)
        ndim = x.ndim
        x = x.reshape((-1, self.dim))

        grads = np.zeros_like(x)

        for i in range(len(grads)):
            xi = x[i]
            grads[i] = numgrad(self.logpdf, xi, h=stepsize)

        grads[np.isinf(grads)] = 0
        grads[np.isnan(grads)] = 0

        if ndim == 0 or (ndim == 1 and self.dim > 1):
            grads = grads[0]
        return grads

    def _to_batch(self, x):
        return {p: x[:, i] for i, p in enumerate(self.parameter_names)}
import logging
import numpy as np
import elfi.methods.mcmc as mcmc
from elfi.classifiers.classifier import Classifier, LogisticRegression
from elfi.loader import get_sub_seed
from elfi.methods.bo.acquisition import LCBSC, AcquisitionBase
from elfi.methods.bo.gpy_regression import GPyRegression
from elfi.methods.bo.utils import CostFunction
from elfi.methods.inference.parameter_inference import ParameterInference
from elfi.methods.posteriors import BOLFIREPosterior
from elfi.methods.results import BOLFIRESample
from elfi.methods.utils import arr2d_to_batch, batch_to_arr2d, resolve_sigmas
from elfi.model.elfi_model import ElfiModel, Summary
from elfi.model.extensions import ModelPrior
               initials=None,
               algorithm='nuts',
               sigma_proposals=None,
               n_evidence=None,
               *args, **kwargs):
        """Sample from the posterior distribution of BOLFIRE.

        Sampling is performed with an MCMC sampler.

        Parameters
        ----------
        n_samples: int
            Number of requested samples from the posterior for each chain. This includes warmup,
            and note that the effective sample size is usually considerably smaller.
        warmup: int, optional
            Length of warmup sequence in MCMC sampling.
        n_chains: int, optional
            Number of independent chains.
        initials: np.ndarray (n_chains, n_params), optional
            Initial values for the sampled parameters for each chain.
        algorithm: str, optional
            Sampling algorithm to use.
        sigma_proposals: np.ndarray
            Standard deviations for Gaussian proposals of each parameter for Metropolis-Hastings.
        n_evidence: int, optional
            If the surrogate model is not fitted yet, specify the amount of evidence.

        Returns
        -------
        BOLFIRESample

        """
        # Fit posterior in case not done
        if self.state['n_batches'] == 0:
            self.fit(n_evidence)

        # Check algorithm
        if algorithm not in ['nuts', 'metropolis']:
            raise ValueError('The given algorithm is not supported.')

        # Check standard deviations of Gaussian proposals when using Metropolis-Hastings
        if algorithm == 'metropolis':
            sigma_proposals = resolve_sigmas(self.target_model.parameter_names,
                                             sigma_proposals,
                                             self.target_model.bounds)

        posterior = self.extract_result()
        warmup = warmup or n_samples // 2

        # Unless given, select the evidence points with best likelihood ratio
        if initials is not None:
            if np.asarray(initials).shape != (n_chains, self.target_model.input_dim):
                raise ValueError('The shape of initials must be (n_chains, n_params).')
        else:
            inds = np.argsort(self.target_model.Y[:, 0])
            initials = np.asarray(self.target_model.X[inds])

        # Enable caching for default RBF kernel
        self.target_model.is_sampling = True

        tasks_ids = []
        ii_initial = 0
        for ii in range(n_chains):
            seed = get_sub_seed(self.seed, ii)
            # Discard bad initialization points
            while np.isinf(posterior.logpdf(initials[ii_initial])):
                ii_initial += 1
                if ii_initial == len(inds):
                    raise ValueError('BOLFIRE.sample: Cannot find enough acceptable '
                                     'initialization points!')

            if algorithm == 'nuts':
                tasks_ids.append(
                    self.client.apply(mcmc.nuts,
                                      n_samples,
                                      initials[ii_initial],
                                      posterior.logpdf,
                                      posterior.gradient_logpdf,
                                      n_adapt=warmup,
                                      seed=seed,
                                      **kwargs))

            elif algorithm == 'metropolis':
                tasks_ids.append(
                    self.client.apply(mcmc.metropolis,
                                      n_samples,
                                      initials[ii_initial],
                                      posterior.logpdf,
                                      sigma_proposals,
                                      warmup,
                                      seed=seed,
                                      **kwargs))

            ii_initial += 1

        # Get results from completed tasks or run sampling (client-specific)
        chains = []
        for id in tasks_ids:
            chains.append(self.client.get_result(id))

        chains = np.asarray(chains)

        logger.info(f'{n_chains} chains of {n_samples} iterations acquired. '
                    'Effective sample size and Rhat for each parameter:')
        for ii, node in enumerate(self.parameter_names):
            logger.info(f'{node} {mcmc.eff_sample_size(chains[:, :, ii])} '
                        f'{mcmc.gelman_rubin_statistic(chains[:, :, ii])}')

        self.target_model.is_sampling = False

        return BOLFIRESample(method_name='BOLFIRE',
                             chains=chains,
                             parameter_names=self.parameter_names,
                             warmup=warmup,
                             n_sim=self.state['n_sim'],
                             seed=self.seed,
                             *args, **kwargs)

    def _resolve_model(self, model):
        """Resolve a given elfi model."""

        if not isinstance(model, ElfiModel):