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

# The six-point example of the lesson, three rounds of boosting with stumps and eta = 0.5
x = np.arange(1, 7, dtype=float)
y = np.array([2, 3, 4, 9, 10, 14], dtype=float)
eta = 0.5


def best_stump(r):
    """The split (midpoint between neighbors) that minimizes the squared error of r."""
    best = None
    for t in (x[:-1] + x[1:]) / 2:
        left, right = r[x < t], r[x > t]
        sse = ((left - left.mean()) ** 2).sum() + ((right - right.mean()) ** 2).sum()
        if best is None or sse < best[0]:
            best = (sse, t, left.mean(), right.mean())
    return best[1:]


F = np.full(6, y.mean())
states, stumps = [F.copy()], []
for m in range(3):
    t, vl, vr = best_stump(y - F)
    stumps.append((t, vl, vr))
    F = F + eta * np.where(x < t, vl, vr)
    states.append(F.copy())

fig, axes = plt.subplots(2, 4, figsize=(10, 4.6), sharex=True,
                         gridspec_kw=dict(height_ratios=[1.25, 1]))
xs = np.linspace(0.5, 6.5, 601)
for m in range(4):
    F = states[m]
    top, bottom = axes[0, m], axes[1, m]
    step = F[np.clip(np.round(xs).astype(int) - 1, 0, 5)]
    top.plot(xs, step, color='#58a6ff', lw=2, label=f'$F_{m}$')
    top.scatter(x, y, color='#f0883e', s=28, zorder=3)
    for xi, yi, fi in zip(x, y, F):
        top.plot([xi, xi], [fi, yi], color='#ff7b72', lw=1, alpha=0.7)
    mse = np.mean((y - F) ** 2)
    top.set_title(f'$F_{m}$' + (' = mean' if m == 0 else f' = $F_{m-1}$ + 0.5·stump {m}')
                  + f'\ntraining MSE {mse:.2f}', fontsize=8.5)
    top.set_ylim(0, 15.5)
    r = y - F
    bottom.axhline(0, color='gray', lw=0.8)
    bottom.vlines(x, 0, r, color='#ff7b72', lw=1.5)
    bottom.scatter(x, r, color='#ff7b72', s=18, zorder=3)
    for xi, ri in zip(x, r):
        bottom.text(xi, ri + (0.5 if ri >= 0 else -0.5), f'{ri:g}', ha='center',
                    va='bottom' if ri >= 0 else 'top', fontsize=7.5, color='gray')
    if m < 3:
        t, vl, vr = stumps[m]
        bottom.plot([0.5, t, t, 6.5], [vl, vl, vr, vr], color='#3fb950', lw=2)
        bottom.set_title(f'stump {m + 1}, fitted to the residuals:\nx < {t:g} → {vl:g}, else {vr:g}',
                         fontsize=7.8, color='#3fb950')
    else:
        bottom.set_title('residuals left\nafter 3 rounds', fontsize=7.8)
    bottom.set_ylim(-7, 9)
    bottom.set_xticks(x)
    for ax in (top, bottom):
        ax.tick_params(labelsize=7.5)
axes[0, 0].set_ylabel('y and prediction', fontsize=8)
axes[1, 0].set_ylabel('residual y − F', fontsize=8)
for ax in axes[1]:
    ax.set_xlabel('x', fontsize=8)

plt.tight_layout()

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

plt.close()