import numpy as np
from tkbo import create_optimizer


def f(x):
    return np.sin(3 * x[:, 0]) * np.exp(-0.1 * x[:, 0])


X = np.linspace(0, 8, 161).reshape(-1, 1)
y = np.full(len(X), np.nan)
idx0 = [0, 80, 160]
y[idx0] = f(X[idx0])

opt = create_optimizer(
    model="sklearn_gpr",
    acquisition="stein",
    acquisition__std_weight=1.0,
    acquisition__grad_weight=1.0,
)

opt.initialize(X, y=y)
res = opt.ask(5)

print("suggested indices:", res.indices)
print("suggested X:", X[res.indices].reshape(-1))