import numpy as np
from tkbo import create_optimizer


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


X = np.linspace(0, 10, 101).reshape(-1, 1)
y = np.full(len(X), np.nan)
idx0 = [0, 25, 50, 75]
y[idx0] = f(X[idx0])

opt = create_optimizer(
    model="sklearn_gpr",
    acquisition="ei",
    surrogate__n_restarts_optimizer=5,
    surrogate__random_state=1,
)

opt.initialize(X, y=y)

for step in range(5):
    res = opt.ask(1)
    idx = int(res.indices[0])
    y_new = f(X[[idx]])[0]
    print(f"step={step} idx={idx} x={X[idx,0]:.3f} y={y_new:.6g} score={res.scores[0]:.6g}")
    opt.tell(idx, y_new)