import matplotlib.pyplot as plt
import numpy as np
from io import StringIO
from matplotlib.colors import LinearSegmentedColormap
from sklearn.datasets import make_moons
from sklearn.tree import DecisionTreeClassifier

# One dataset (200 points), 100 bootstrap resamples of it (200 rows drawn with
# replacement), one fully grown tree per resample. Panels 1-3: three of those
# trees, with their root split drawn thick. Panel 4: the share of the 100
# trees that vote "blue" at each point of the plane.
X, y = make_moons(n_samples=200, noise=0.3, random_state=0)
rng = np.random.default_rng(0)
pad = 0.5
x0, x1 = X[:, 0].min() - pad, X[:, 0].max() + pad
y0, y1 = X[:, 1].min() - pad, X[:, 1].max() + pad
xx, yy = np.meshgrid(np.linspace(x0, x1, 250), np.linspace(y0, y1, 250))
grid = np.c_[xx.ravel(), yy.ravel()]

trees, votes = [], []
for b in range(100):
    idx = rng.integers(0, len(X), len(X))
    t = DecisionTreeClassifier(random_state=0).fit(X[idx], y[idx])
    trees.append((t, idx))
    votes.append(t.predict(grid))
votes = np.array(votes)
mean_vote = votes.mean(axis=0)
# how often two trees disagree, averaged over all pairs and over the plane
agree_frac = mean_vote ** 2 + (1 - mean_vote) ** 2
n = len(votes)
pair_dis = np.mean(1 - (n * agree_frac - 1) / (n - 1))

BLUE, ORANGE = '#58a6ff', '#f0883e'
plt.rcParams.update({'text.color': 'gray', 'axes.labelcolor': 'gray', 'axes.edgecolor': 'gray',
                     'xtick.color': 'gray', 'ytick.color': 'gray'})
fig, axes = plt.subplots(1, 4, figsize=(10.5, 3.0))
names = ['x₁', 'x₂']
for k, ax in enumerate(axes[:3]):
    t, idx = trees[k]
    Z = t.predict(grid).reshape(xx.shape)
    ax.contourf(xx, yy, Z, levels=[-0.5, 0.5, 1.5], colors=[ORANGE, BLUE], alpha=0.2).set_rasterized(True)
    ax.scatter(X[idx][y[idx] == 0, 0], X[idx][y[idx] == 0, 1], s=6, color=ORANGE)
    ax.scatter(X[idx][y[idx] == 1, 0], X[idx][y[idx] == 1, 1], s=6, color=BLUE)
    f, thr = t.tree_.feature[0], t.tree_.threshold[0]
    if f == 0:
        ax.axvline(thr, color='gray', lw=2)
    else:
        ax.axhline(thr, color='gray', lw=2)
    ax.set_title(f'resample {k + 1}: root {names[f]} ≤ {thr:.2f}\n{t.get_n_leaves()} leaves', fontsize=8.5)
ax = axes[3]
cmap = LinearSegmentedColormap.from_list('vote', [ORANGE, 'white', BLUE])
im = ax.contourf(xx, yy, mean_vote.reshape(xx.shape), levels=np.linspace(0, 1, 11), cmap=cmap, alpha=0.45)
im.set_rasterized(True)
ax.scatter(X[y == 0, 0], X[y == 0, 1], s=6, color=ORANGE)
ax.scatter(X[y == 1, 0], X[y == 1, 1], s=6, color=BLUE)
ax.set_title(f'all 100 trees: share voting blue\ntwo trees disagree on {pair_dis:.1%} of the plane', fontsize=8.5)
for ax in axes:
    ax.set_xlim(x0, x1); ax.set_ylim(y0, y1)
    ax.set_xticks([]); ax.set_yticks([])
plt.tight_layout()

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

plt.close()