import warnings
import matplotlib.pyplot as plt
import numpy as np
from io import StringIO
from sklearn.datasets import load_digits
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import StratifiedKFold, cross_val_score, train_test_split

warnings.filterwarnings('ignore')          # with very few trees some rows have no OOB vote yet

# Grow one forest on the digits, tree by tree (warm_start), and after each size record the
# out-of-bag error, the error on the 450 test images and, at a few sizes, the 5-fold CV error
X, y = load_digits(return_X_y=True)
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.25, stratify=y, random_state=0)
cv = StratifiedKFold(n_splits=5, shuffle=True, random_state=0)

sizes = [5, 8, 12, 20, 30, 50, 80, 120, 200, 300, 500, 800, 1200, 2000]
forest = RandomForestClassifier(warm_start=True, oob_score=True, random_state=0, n_jobs=-1)
oob, test = [], []
for B in sizes:
    forest.set_params(n_estimators=B).fit(X_train, y_train)
    proba = forest.oob_decision_function_
    seen = ~np.isnan(proba).any(axis=1)                        # rows with at least one OOB tree
    oob.append(np.mean(forest.classes_[proba[seen].argmax(axis=1)] != y_train[seen]))
    test.append(1 - forest.score(X_test, y_test))
cv_sizes = [10, 30, 100, 300, 1000]
cv_err = [1 - cross_val_score(RandomForestClassifier(n_estimators=B, random_state=0, n_jobs=-1),
                              X_train, y_train, cv=cv).mean() for B in cv_sizes]

fig, ax = plt.subplots(figsize=(7.5, 3.5))
ax.plot(sizes, oob, 'o-', color='#58a6ff', ms=4, label='out-of-bag error (training rows, free)')
ax.plot(sizes, test, 's-', color='#f0883e', ms=4, label='test error (450 untouched images)')
ax.plot(cv_sizes, cv_err, 'D', color='#3fb950', ms=5, label='5-fold CV error (5 extra fits)')
ax.annotate(f'5 trees: OOB {oob[0]:.3f} vs test {test[0]:.3f}\n(each row is voted on by only ~2 trees)',
            (sizes[0], oob[0]), xytext=(9, 0.2), fontsize=7.5, arrowprops=dict(arrowstyle='->', color='gray'))
ax.text(2000, oob[-1] * 1.25, f'2000 trees: OOB {oob[-1]:.3f} · test {test[-1]:.3f}',
        fontsize=7.5, ha='right', color='gray')
ax.set_xscale('log')
ax.set_xticks([5, 10, 30, 100, 300, 1000, 2000], ['5', '10', '30', '100', '300', '1000', '2000'])
ax.set_xlabel('number of trees (n_estimators)')
ax.set_ylabel('error (log scale)')
ax.set_yscale('log')
ax.set_ylim(0.012, 0.35)
ax.set_yticks([0.02, 0.05, 0.1, 0.2], ['0.02', '0.05', '0.10', '0.20'])
ax.minorticks_off()
ax.legend(fontsize=7.5, loc='upper right')
ax.set_title('More trees: the error falls, then flattens — it never turns back up', fontsize=9.5, loc='left')
ax.tick_params(labelsize=8)

plt.tight_layout()

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

plt.close()