import matplotlib.pyplot as plt
from io import StringIO
from matplotlib.patches import FancyBboxPatch

# The boosting loop as a cycle, with the numbers of round 1 of the lesson's six-point example
# (x = 1..6, y = 2, 3, 4, 9, 10, 14, squared loss, learning rate 0.5).
BLUE, ORANGE, GREEN, PURPLE = '#58a6ff', '#f0883e', '#3fb950', '#bc8cff'
plt.rcParams.update({'text.color': 'gray'})
fig, ax = plt.subplots(figsize=(10, 6.6))
ax.set_xlim(0, 10); ax.set_ylim(0, 6.6); ax.axis('off')


def box(x, y, title, body, color, w=3.5, h=1.25):
    ax.add_patch(FancyBboxPatch((x - w / 2, y - h / 2), w, h, boxstyle='round,pad=0.04,rounding_size=0.15',
                                fc=color, ec=color, alpha=0.16, lw=1.5))
    ax.add_patch(FancyBboxPatch((x - w / 2, y - h / 2), w, h, boxstyle='round,pad=0.04,rounding_size=0.15',
                                fc='none', ec=color, lw=1.5))
    ax.text(x, y + 0.3, title, ha='center', va='center', fontsize=10.5, fontweight='bold', color=color)
    ax.text(x, y - 0.18, body, ha='center', va='center', fontsize=9.3, linespacing=1.4)


def arrow(xa, ya, xb, yb, rad=0.0):
    ax.annotate('', xy=(xb, yb), xytext=(xa, ya),
                arrowprops=dict(arrowstyle='-|>', color='gray', lw=1.6, connectionstyle=f'arc3,rad={rad}'))


box(5, 6.05, 'start: F₀ = the constant that minimizes the loss', 'squared loss → the mean: 42 / 6 = 7', GREEN,
    w=5.6, h=0.9)
box(2.3, 4.25, '1 · current predictions', 'F₀ = 7, 7, 7, 7, 7, 7', BLUE)
box(7.7, 4.25, '2 · pseudo-residuals', 'r = y − F = −5, −4, −3, 2, 3, 7\n(the negative gradient of the loss)', PURPLE)
box(7.7, 1.45, '3 · fit a small tree to r', 'best stump: x < 3.5\nleaves −4 / +4', ORANGE)
box(2.3, 1.45, '4 · add it, shrunk by η', 'F₁ = F₀ + 0.5 · stump\n= 5, 5, 5, 9, 9, 9', GREEN)
arrow(5, 5.6, 2.9, 4.9)
arrow(4.1, 4.3, 5.9, 4.3)
arrow(7.7, 3.6, 7.7, 2.1)
arrow(5.9, 1.45, 4.1, 1.45)
arrow(2.3, 2.1, 2.3, 3.6)
ax.text(5, 2.88, 'repeat M times', ha='center', fontsize=11, fontweight='bold')
ax.text(5, 2.45, 'training MSE: 18.67 → 6.67 → 2.92 → 1.23', ha='center', fontsize=9.5)
ax.text(5, 0.25, 'the tree turns the residuals, known only at the training points, into a function of x '
        '— and η makes every step a small one', ha='center', fontsize=9, style='italic')

buffer = StringIO()
plt.savefig(buffer, format="svg", transparent=True, bbox_inches='tight', pad_inches=0.05)
print(buffer.getvalue())

plt.close()