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

# What max_features does at a split. In bagging every feature is a candidate at every node, so the
# strongest one (f3 here) wins the root of almost every tree and the trees look alike. In a random
# forest each node draws its own small set of candidates; when f3 is not drawn, another feature
# gets the split, and the trees diverge.
BLUE, ORANGE, GREEN, PURPLE, GRAY = '#58a6ff', '#f0883e', '#3fb950', '#bc8cff', '#8b949e'
plt.rcParams.update({'text.color': 'gray'})
fig, axes = plt.subplots(1, 2, figsize=(11, 5.0))
d = 8
strength = [0.25, 0.4, 0.3, 0.95, 0.35, 0.5, 0.2, 0.45]          # how good each feature's best split is


def chips(ax, y, candidates, chosen, label):
    ax.text(-0.3, y + 0.35, label, ha='right', va='center', fontsize=9)
    for j in range(d):
        cand = j in candidates
        fc = PURPLE if j == chosen else (BLUE if cand else GRAY)
        ax.add_patch(FancyBboxPatch((j * 0.62, y), 0.52, 0.7, boxstyle='round,pad=0,rounding_size=0.08',
                                    fc=fc, ec='none', alpha=0.85 if j == chosen else (0.35 if cand else 0.12)))
        ax.text(j * 0.62 + 0.26, y + 0.35, f'f{j}', ha='center', va='center', fontsize=8.5,
                fontweight='bold' if j == chosen else 'normal')


rng = np.random.default_rng(2)
for ax, (title, m) in zip(axes, [('bagging: max_features = all 8', d), ('random forest: max_features = 3', 3)]):
    ax.set_xlim(-2.6, d * 0.62 + 0.3); ax.set_ylim(-1.4, 5.6); ax.axis('off')
    ax.set_title(title, fontsize=11, fontweight='bold', color='gray')
    roots = []
    for k in range(5):
        cand = list(range(d)) if m == d else sorted(rng.choice(d, m, replace=False))
        chosen = max(cand, key=lambda j: strength[j])
        roots.append(chosen)
        chips(ax, 4.4 - k * 1.0, cand, chosen, f'tree {k + 1}, root')
    same = sum(r == 3 for r in roots)
    msg = (f'f3 is the best split everywhere:\n{same} of 5 roots are f3 → similar trees,\ncorrelated errors'
           if m == d else
           f'f3 is only a candidate when drawn:\n{same} of 5 roots are f3 → different trees,\nless correlated errors')
    ax.text(d * 0.31, -1.0, msg, ha='center', fontsize=9.5, color=PURPLE if m < d else ORANGE, fontweight='bold')
fig.text(0.5, 0.005, 'purple: the candidate the node splits on · blue: drawn as candidates at this node · gray: not '
         'eligible here (drawn again at every node)', ha='center', fontsize=8.5, style='italic', color='gray')
plt.tight_layout(rect=(0, 0.04, 1, 1))

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

plt.close()