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

# Bootstrap and out-of-bag, as one table. Each column is one tree's bootstrap sample of the same
# 10 training rows: the number is how many times the row was drawn, an empty gray cell means the
# row was left out (out-of-bag) for that tree. The highlighted row is predicted out-of-bag only
# by the trees that never saw it.
BLUE, ORANGE, GREEN, PURPLE = '#58a6ff', '#f0883e', '#3fb950', '#bc8cff'
n_rows, n_trees, row = 10, 6, 6
rng = np.random.default_rng(4)
counts = np.array([np.bincount(rng.integers(0, n_rows, n_rows), minlength=n_rows) for _ in range(n_trees)]).T
oob_trees = np.flatnonzero(counts[row] == 0)

plt.rcParams.update({'text.color': 'gray'})
fig, ax = plt.subplots(figsize=(10.5, 5.6))
ax.set_xlim(-2.2, n_trees + 4.2); ax.set_ylim(-2.3, n_rows + 1.6); ax.axis('off')

for j in range(n_trees):
    ax.text(j + 0.5, n_rows + 0.9, f'tree {j + 1}', ha='center', fontsize=9.5, fontweight='bold')
    ax.text(j + 0.5, n_rows + 0.35, f'{(counts[:, j] == 0).sum()} left out', ha='center', fontsize=8)
for i in range(n_rows):
    y = n_rows - 1 - i
    ax.text(-0.25, y + 0.5, f'row {i}', ha='right', va='center', fontsize=9,
            fontweight='bold' if i == row else 'normal', color=PURPLE if i == row else 'gray')
    for j in range(n_trees):
        c = counts[i, j]
        fc = 'gray' if c == 0 else BLUE
        alpha = 0.12 if c == 0 else 0.15 + 0.15 * c
        ax.add_patch(FancyBboxPatch((j + 0.06, y + 0.06), 0.88, 0.88, boxstyle='round,pad=0,rounding_size=0.08',
                                    fc=fc, ec='none', alpha=alpha))
        if c:
            ax.text(j + 0.5, y + 0.5, f'×{c}', ha='center', va='center', fontsize=9)
        else:
            ax.text(j + 0.5, y + 0.5, 'OOB', ha='center', va='center', fontsize=7.5, style='italic')
    if i == row:
        ax.add_patch(FancyBboxPatch((0.0, y), n_trees, 1, boxstyle='round,pad=0.02', fc='none', ec=PURPLE, lw=2))

names = ', '.join(str(j + 1) for j in oob_trees)
ax.annotate(f'row {row} was never drawn by\ntrees {names}: only they vote\non it → its out-of-bag prediction',
            xy=(n_trees + 0.05, n_rows - 1 - row + 0.5), xytext=(n_trees + 0.6, n_rows - 1 - row + 2.6),
            fontsize=9.5, color=PURPLE, fontweight='bold',
            arrowprops=dict(arrowstyle='->', color=PURPLE, lw=1.3))
ax.text(n_trees + 0.6, 2.6, 'each column sums to 10:\n10 draws with replacement,\nsome rows twice or more,\nabout 37% never', fontsize=9,
        va='top')
ax.text(n_trees / 2, -0.7, 'blue = drawn into that tree\'s bootstrap sample (darker = more copies) · gray = out-of-bag',
        ha='center', fontsize=9, style='italic')
ax.text(n_trees / 2, -1.5, f'on average {(counts == 0).mean():.0%} of the cells here are out-of-bag; '
        'with many rows the share tends to 1/e ≈ 36.8%', 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()