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

# The two tree ensembles side by side. A random forest grows deep trees independently (in parallel),
# each on its own bootstrap sample, and averages them: it fights variance. Gradient boosting grows
# small trees one after another, each fitted to what the ensemble still gets wrong, and adds them up
# with a learning rate: it fights bias.
BLUE, ORANGE, GREEN, PURPLE = '#58a6ff', '#f0883e', '#3fb950', '#bc8cff'
plt.rcParams.update({'text.color': 'gray'})
fig, ax = plt.subplots(figsize=(11, 6.6))
ax.set_xlim(0, 11); ax.set_ylim(0, 6.6); ax.axis('off')


def tree(x, y, depth, color, scale=0.32):
    """A small tree icon: depth levels of nodes, drawn top-down from (x, y)."""
    levels = [[(x, y)]]
    for d in range(1, depth + 1):
        w = scale * 2 ** (depth - d) * 0.9
        levels.append([(px + s * w / 2, y - d * scale) for px, _ in levels[-1] for s in (-1, 1)])
    for d in range(1, len(levels)):
        for k, (cx, cy) in enumerate(levels[d]):
            px, py = levels[d - 1][k // 2]
            ax.plot([px, cx], [py, cy], color=color, lw=1.1, alpha=0.8)
    for lvl in levels:
        for cx, cy in lvl:
            ax.plot(cx, cy, 'o', ms=4.5, color=color)


def box(x, y, text, color, w=1.5, h=0.55):
    ax.add_patch(FancyBboxPatch((x - w / 2, y - h / 2), w, h, boxstyle='round,pad=0.03,rounding_size=0.1',
                                fc=color, ec=color, alpha=0.18))
    ax.text(x, y, text, ha='center', va='center', fontsize=9.5, fontweight='bold', color='gray')


def arrow(xa, ya, xb, yb):
    ax.annotate('', xy=(xb, yb), xytext=(xa, ya), arrowprops=dict(arrowstyle='-|>', color='gray', lw=1.1))


# --- random forest (top) ---
ax.text(0.1, 6.35, 'RANDOM FOREST · bagging', fontsize=11, fontweight='bold', color=GREEN)
ax.text(0.1, 6.0, 'deep trees, grown independently (in parallel), then averaged', fontsize=9.5)
box(1.0, 5.0, 'training data', GREEN, w=1.5)
xs = [3.2, 4.9, 6.6, 8.3]
for k, x in enumerate(xs):
    ax.text(x, 5.6, f'bootstrap {k + 1}', ha='center', fontsize=8.5)
    ax.annotate('', xy=(x - 0.1, 5.45), xytext=(1.8, 5.1),
                arrowprops=dict(arrowstyle='-|>', color='gray', lw=1.0, connectionstyle=f'arc3,rad=-{0.08 + 0.03 * k}'))
    tree(x, 5.35, 3, GREEN, scale=0.2)
    arrow(x, 4.7, 5.75 + (x - 5.75) * 0.25, 4.15)
box(5.75, 3.9, 'average / vote', GREEN, w=1.9, h=0.45)
ax.text(5.75, 3.35, 'each tree: low bias, high variance  →  the average cancels the variance', ha='center',
        fontsize=9, style='italic')

# --- boosting (bottom) ---
ax.text(0.1, 2.95, 'GRADIENT BOOSTING', fontsize=11, fontweight='bold', color=ORANGE)
ax.text(0.1, 2.62, 'small trees, one after another, each fitted to the residuals of the sum so far', fontsize=9.5)
Y = 0.95
box(0.95, Y, 'F₀ = mean', ORANGE, w=1.3, h=0.45)
xs = [2.85, 4.95, 7.05]
for k, x in enumerate(xs):
    arrow(1.6 if k == 0 else xs[k - 1] + 0.95, Y, x - 0.95, Y)
    ax.text(x, 2.05, f'h{k + 1}: fitted to the\nresiduals of F{k}', ha='center', fontsize=8.3)
    tree(x, 1.8, 1, ORANGE, scale=0.3)
    arrow(x, 1.45, x, Y + 0.25)
    box(x, Y, f'F{k + 1} = F{k} + η·h{k + 1}', ORANGE, w=1.85, h=0.45)
ax.text(8.45, Y, '· · ·', fontsize=14, ha='center', va='center')
arrow(8.75, Y, 9.35, Y)
box(9.85, Y, 'F_M', ORANGE, w=0.9, h=0.45)
ax.text(5.4, 0.25, 'each tree: high bias, low variance  →  the sum removes the bias step by step', 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()