import matplotlib.pyplot as plt
import numpy as np
from io import StringIO
from scipy.stats import norm
from sklearn.gaussian_process import GaussianProcessRegressor
from sklearn.gaussian_process.kernels import RBF


# A stand-in for an expensive validation score, as a function of one (log-scaled) hyperparameter:
# a small peak at 0.25 and a higher, narrower one at 0.72.
def score(x):
    return 0.70 + 0.08 * np.exp(-((x - 0.25) / 0.10) ** 2) + 0.12 * np.exp(-((x - 0.72) / 0.06) ** 2)


xs = np.linspace(0, 1, 201)
X = np.array([0.15, 0.35, 0.95])                 # three evaluations to start
Y = score(X)

fig, axes = plt.subplots(2, 4, figsize=(10, 4.6), sharex=True, gridspec_kw={'height_ratios': [2.2, 1]})
for it in range(4):
    gp = GaussianProcessRegressor(kernel=RBF(length_scale=0.15), normalize_y=True, optimizer=None)
    gp.fit(X.reshape(-1, 1), Y)
    mu, sd = gp.predict(xs.reshape(-1, 1), return_std=True)
    best = Y.max()
    z = (mu - best) / np.maximum(sd, 1e-12)
    ei = (mu - best) * norm.cdf(z) + sd * norm.pdf(z)
    x_next = xs[ei.argmax()]

    top, bot = axes[0, it], axes[1, it]
    top.plot(xs, score(xs), color='gray', lw=1, ls='--', label='true score (unknown)')
    top.fill_between(xs, mu - 2 * sd, mu + 2 * sd, color='#58a6ff', alpha=0.18, lw=0)
    top.plot(xs, mu, color='#58a6ff', lw=1.6, label='surrogate mean ± 2 sd')
    top.scatter(X, Y, color='#f0883e', s=26, zorder=3, label='evaluated')
    top.axhline(best, color='#3fb950', lw=0.8, ls=':')
    top.axvline(x_next, color='#ff7b72', lw=1)
    top.set_ylim(0.6, 0.9)
    top.set_title(f'iteration {it + 1}: {len(X)} evaluations, best {best:.3f}', fontsize=8.5)
    bot.fill_between(xs, 0, ei, color='#ff7b72', alpha=0.35, lw=0)
    bot.plot(xs, ei, color='#ff7b72', lw=1.2)
    bot.scatter([x_next], [ei.max()], color='#ff7b72', s=20, zorder=3)
    bot.annotate(f'next: x = {x_next:.3f}', (x_next, ei.max()), textcoords='offset points',
                 xytext=(-30 if x_next > 0.6 else 4, -2), fontsize=7, color='#ff7b72', ha='right' if x_next > 0.6 else 'left')
    bot.set_ylim(0, ei.max() * 1.35)
    bot.set_yticks([])
    bot.set_xlabel('hyperparameter x', fontsize=8)
    for ax in (top, bot):
        ax.tick_params(labelsize=7)
    if it == 0:
        top.set_ylabel('validation score', fontsize=8)
        bot.set_ylabel('expected\nimprovement', fontsize=8)
        top.legend(fontsize=6.5, loc='upper left', frameon=False)
    X = np.append(X, x_next)
    Y = np.append(Y, score(x_next))

fig.suptitle('Bayesian optimization: fit the surrogate, evaluate where the expected improvement is largest',
             fontsize=9.5)
plt.tight_layout()

buffer = StringIO()
plt.savefig(buffer, format="svg", transparent=True)
print(buffer.getvalue())

plt.close()