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.ensemble import RandomForestClassifier
from sklearn.tree import DecisionTreeClassifier

# 300 noisy "moons" to train on, 10,000 fresh points from the same process to score on
X, y = make_moons(n_samples=300, noise=0.3, random_state=0)
X_new, y_new = make_moons(n_samples=10_000, noise=0.3, random_state=1)

tree = DecisionTreeClassifier(random_state=0).fit(X, y)
forest = RandomForestClassifier(n_estimators=300, random_state=0).fit(X, y)
tree_accs = [t.score(X_new, y_new) for t in forest.estimators_]

xx, yy = np.meshgrid(np.linspace(-2, 3, 250), np.linspace(-1.6, 2.1, 200))
grid = np.c_[xx.ravel(), yy.ravel()]


def dots(ax):
    ax.scatter(X[y == 0, 0], X[y == 0, 1], c='#f0883e', s=9, lw=0)
    ax.scatter(X[y == 1, 0], X[y == 1, 1], c='#58a6ff', s=9, lw=0)
    ax.set_xticks([]); ax.set_yticks([])


fig, axes = plt.subplots(1, 3, figsize=(10, 3.3))

ax = axes[0]
Z = tree.predict_proba(grid)[:, 1].reshape(xx.shape)
ax.contourf(xx, yy, Z, levels=[-0.01, 0.5, 1.01], colors=['#f0883e', '#58a6ff'], alpha=0.25)
ax.contour(xx, yy, Z, levels=[0.5], colors='gray', linewidths=0.8)
dots(ax)
ax.set_title(f'one full tree\naccuracy on new data {tree.score(X_new, y_new):.3f}', fontsize=9)

ax = axes[1]
for t, col in zip(forest.estimators_[:3], ['#3fb950', '#ff7b72', 'gray']):
    Zt = t.predict_proba(grid)[:, 1].reshape(xx.shape)
    ax.contour(xx, yy, Zt, levels=[0.5], colors=col, linewidths=1.3)
dots(ax)
ax.set_title(f'3 of its 300 trees, each on its own bootstrap\n'
             f'each tree alone: accuracy {np.mean(tree_accs):.3f} on average', fontsize=9)

ax = axes[2]
Z = forest.predict_proba(grid)[:, 1].reshape(xx.shape)
cs = ax.contourf(xx, yy, Z, levels=np.linspace(0, 1, 11), cmap=LinearSegmentedColormap.from_list('ob', ['#f0883e', '#ffffff', '#58a6ff']), alpha=0.4)
for c in cs.collections if hasattr(cs, 'collections') else [cs]:
    c.set_rasterized(True)
ax.contour(xx, yy, Z, levels=[0.5], colors='gray', linewidths=1.2)
dots(ax)
ax.set_title(f'random forest: the vote of the 300 trees\naccuracy on new data {forest.score(X_new, y_new):.3f}',
             fontsize=9)

plt.tight_layout()

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

plt.close()