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, train_test_split

# The lesson's max_features experiment on the digits, for every value from 1 to 64 pixels:
# 5-fold CV on the training part, 200 trees per forest
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)

values = [1, 2, 3, 4, 6, 8, 12, 16, 24, 32, 48, 64]
corr, tree_err, forest_err = [], [], []
for m in values:
    c, te, fe = [], [], []
    for tr, va in cv.split(X_train, y_train):
        rf = RandomForestClassifier(n_estimators=200, max_features=m, random_state=0, n_jobs=-1)
        rf.fit(X_train[tr], y_train[tr])
        Xv, yv = X_train[va], y_train[va]
        wrong = np.array([rf.classes_[t.predict(Xv).astype(int)] != yv for t in rf.estimators_])
        C = np.corrcoef(wrong)
        c.append(np.nanmean(C[np.triu_indices(200, 1)]))
        te.append(wrong.mean())
        fe.append(1 - rf.score(Xv, yv))
    corr.append(np.mean(c)); tree_err.append(np.mean(te)); forest_err.append(np.mean(fe))

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(9.5, 3.5))
ax1.plot(values, corr, 'o-', color='#58a6ff', ms=4, label='correlation between the mistakes of two trees')
ax1.plot(values, tree_err, 'o-', color='#f0883e', ms=4, label='error of one tree (5-fold CV)')
ax1.set_title('Fewer candidates per split: trees more different,\nand each one worse', fontsize=9, loc='left')
ax1.set_ylim(0, 0.42)
ax1.legend(fontsize=7.5, loc='upper center')

ax2.plot(values, forest_err, 'o-', color='#3fb950', ms=4)
best = int(np.argmin(forest_err))
ax2.annotate(f'best: max_features = {values[best]}\nerror {forest_err[best]:.3f}',
             (values[best], forest_err[best]), xytext=(2.2, 0.045), fontsize=7.5,
             arrowprops=dict(arrowstyle='->', color='gray'))
ax2.annotate(f'bagging (all 64)\nerror {forest_err[-1]:.3f}', (64, forest_err[-1]), xytext=(20, 0.058),
             fontsize=7.5, arrowprops=dict(arrowstyle='->', color='gray'))
ax2.annotate(f'1 pixel\nerror {forest_err[0]:.3f}', (1, forest_err[0]), xytext=(1.05, 0.012),
             fontsize=7.5, arrowprops=dict(arrowstyle='->', color='gray'))
ax2.set_ylabel('error of the forest (5-fold CV)', fontsize=8)
ax2.set_title('The forest of 200 trees: best in between', fontsize=9, loc='left')
ax2.set_ylim(0, 0.07)
for ax in (ax1, ax2):
    ax.set_xscale('log')
    ax.set_xticks([1, 2, 4, 8, 16, 32, 64], ['1', '2', '4', '8\n√64', '16', '32', '64\nall'])
    ax.set_xlabel('max_features (pixels considered at each split)', fontsize=8)
    ax.tick_params(labelsize=8)
    ax.axvline(8, color='gray', ls=':', lw=1)

plt.tight_layout()

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

plt.close()