import tensorflow as tf


search_space = Box([0, 0], [1, 1])
num_initial_points = 3
initial_query_points = search_space.sample(num_initial_points)
initial_observations = objective(initial_query_points.numpy(), sleep=False)
initial_data = Dataset(
    query_points=initial_query_points,
    observations=tf.constant(initial_observations, dtype=tf.float64),
)

import gpflow
from trieste.models.gpflow import GaussianProcessRegression


def build_model(data):
    variance = tf.math.reduce_variance(data.observations)
    kernel = gpflow.kernels.RBF(variance=variance)
    gpr = gpflow.models.GPR(data.astuple(), kernel, noise_variance=1e-5)
    gpflow.set_trainable(gpr.likelihood, False)
    return GaussianProcessRegression(gpr)

# these imports will be used later for optimization
from trieste.acquisition import LocalPenalizationAcquisitionFunction
from trieste.acquisition.rule import AsynchronousGreedy, EfficientGlobalOptimization
from trieste.ask_tell_optimization import AskTellOptimizer


# %% [markdown]
# ## Multiprocessing setup
#
