import matplotlib.pyplot as plt
import numpy as np
from io import StringIO

# Bergstra & Bengio's argument on a synthetic objective: the score depends on ONE
# hyperparameter (x1, a bump of width 0.05 around an unknown optimum c); the other
# d - 1 hyperparameters do nothing. Grid and random search get the same budget B.
W = 0.05


def best_score(x1_tried, c):
    """Mean best score over problems c; x1_tried is (B,) for a grid, (len(c), B) for random."""
    x = np.broadcast_to(x1_tried, (len(c), np.shape(x1_tried)[-1]))
    return np.exp(-(x - c[:, None]) ** 2 / (2 * W ** 2)).max(axis=1).mean()


rng = np.random.default_rng(0)
c = rng.uniform(0, 1, 4000)                      # 4000 problems, optimum anywhere

fig, (left, right) = plt.subplots(1, 2, figsize=(10, 4.6), gridspec_kw={'width_ratios': [1, 1.25]})

# Left: 9 evaluations in 2 dimensions, projected on the important axis
c0 = 0.62
k = 3
g = (np.arange(k) + 0.5) / k
gx, gy = np.meshgrid(g, g)
rx, ry = rng.uniform(0, 1, 9), rng.uniform(0, 1, 9)
xs = np.linspace(0, 1, 400)
bump = np.exp(-(xs - c0) ** 2 / (2 * W ** 2))
left.plot(xs, 1.08 + 0.35 * bump, color='gray', lw=1.4)
left.text(0.02, 1.36, 'score (depends on x1 only)', fontsize=7.5, color='gray')
left.scatter(gx.ravel(), gy.ravel(), s=46, color='#f0883e', label='grid 3 × 3: 3 distinct x1 values')
left.scatter(rx, ry, s=46, color='#58a6ff', marker='D', label='random: 9 distinct x1 values')
for x in g:
    left.plot([x] * 2, [1.03, 1.08], color='#f0883e', lw=2.2)
for x in rx:
    left.plot([x] * 2, [1.00, 1.05], color='#58a6ff', lw=1.6)
left.axvline(c0, color='#3fb950', lw=1, ls='--')
left.text(c0 + 0.015, 0.03, 'best x1', color='#3fb950', fontsize=7.5)
left.set_xlim(0, 1)
left.set_ylim(0, 1.48)
left.set_yticks([0, 0.5, 1])
left.set_xlabel('x1: the hyperparameter that matters', fontsize=8.5)
left.set_ylabel('x2: a hyperparameter that does not', fontsize=8.5)
left.set_title('Same budget, 9 evaluations', fontsize=9.5)
left.legend(fontsize=7.5, loc='upper center', bbox_to_anchor=(0.5, -0.16), ncol=2, frameon=False)
left.tick_params(labelsize=7.5)

# Right: mean best score against the budget
budgets = np.arange(4, 257)
rand = [best_score(rng.uniform(0, 1, (len(c), B)), c) for B in budgets]   # fresh points per problem
right.plot(budgets, rand, color='#58a6ff', lw=1.8, label='random (any number of dimensions)')
for d, col, mk in [(2, '#f0883e', 'o'), (3, '#ff7b72', 's'), (4, 'gray', '^')]:
    ks = np.arange(2, 17)
    ks = ks[ks ** d <= 256]
    Bs = ks ** d
    sc = [best_score((np.arange(kk) + 0.5) / kk, c) for kk in ks]
    right.plot(Bs, sc, color=col, marker=mk, ms=5, lw=1.2, label=f'grid, {d} hyperparameters')
    for B, s, kk in zip(Bs, sc, ks):
        if (d, B) in [(2, 9), (2, 81), (3, 27), (3, 125), (4, 16), (4, 81), (4, 256)]:
            right.annotate(f'{kk} values', (B, s), textcoords='offset points', xytext=(5, -11),
                           fontsize=7, color=col)
right.set_xscale('log')
right.set_xticks([4, 9, 16, 27, 64, 81, 256])
right.set_xticklabels(['4', '9', '16', '27', '64', '81', '256'])
right.set_ylim(0, 1.02)
right.set_xlabel('budget: number of configurations evaluated (log scale)', fontsize=8.5)
right.set_ylabel('best score found (1 = the optimum)', fontsize=8.5)
right.set_title('Mean over 4,000 problems: the grid wastes its budget on x2, x3, ...', fontsize=9.5)
right.legend(fontsize=7.5, loc='lower right')
right.grid(alpha=0.25)
right.tick_params(labelsize=7.5)

plt.tight_layout()

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

plt.close()