import pandas as pd

from baybe.acquisition import ProbabilityOfImprovement
from baybe.objectives import DesirabilityObjective
from baybe.parameters import NumericalContinuousParameter
from baybe.recommenders import BotorchRecommender
from baybe.targets import NumericalTarget

searchspace = NumericalContinuousParameter("p", [0, 1]).to_searchspace()
targets = [NumericalTarget("t1"), NumericalTarget("t2")]
objective = DesirabilityObjective(
    targets, scalarizer="MEAN", require_normalization=False
)
recommender = BotorchRecommender()

candidates = pd.DataFrame({"p": [0.2, 0.3]})
measurements = pd.DataFrame(
    {
        "p": [0.1, 0.4, 0.6, 0.8],
        "t1": [0.2, 0.5, 0.7, 0.9],
        "t2": [0.3, 0.6, 0.4, 0.8],
    }
)

acqf = recommender.acquisition_values(
    candidates,
    searchspace,
    objective,
    measurements,
    acquisition_function=ProbabilityOfImprovement(),
)
print(acqf)