import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from io import StringIO
from sklearn.datasets import fetch_openml
from sklearn.model_selection import GridSearchCV, StratifiedKFold, train_test_split
from sklearn.tree import DecisionTreeClassifier

# Spambase without its repeated rows, the development part of the lesson's
# split. Left: pre-pruning by max_depth. Right: post-pruning by ccp_alpha,
# plotted against the number of leaves of the pruned tree. Both panels use the
# lesson's 5-fold CV; the green ring marks the one-standard-error choice.
spam = fetch_openml(data_id=44, as_frame=True)
X, y = spam.data, spam.target.astype(int)
keep = ~X.duplicated()
X, y = X[keep], y[keep]
X_dev, X_test, y_dev, 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)

BLUE, ORANGE, GREEN = '#58a6ff', '#f0883e', '#3fb950'
plt.rcParams.update({'text.color': 'gray', 'axes.labelcolor': 'gray', 'axes.edgecolor': 'gray',
                     'xtick.color': 'gray', 'ytick.color': 'gray'})
fig, axes = plt.subplots(1, 2, figsize=(10, 3.9), sharey=True)


def curves(ax, grid_values, param, xvals, xlabel):
    g = GridSearchCV(DecisionTreeClassifier(random_state=0), {param: grid_values}, cv=cv,
                     return_train_score=True).fit(X_dev, y_dev)
    r = pd.DataFrame(g.cv_results_)
    tr, va, sd = r['mean_train_score'], r['mean_test_score'], r['std_test_score']
    ax.fill_between(xvals, va - sd, va + sd, color=BLUE, alpha=0.15, lw=0)
    ax.plot(xvals, tr, 'o-', color=ORANGE, ms=3.5, label='training folds')
    ax.plot(xvals, va, 'o-', color=BLUE, ms=3.5, label='validation folds (± 1 std)')
    b = int(va.idxmax())
    ax.scatter([xvals[b]], [va[b]], s=90, marker='*', color=BLUE, zorder=6)
    se = sd[b] / np.sqrt(5)
    ok = np.where(va >= va[b] - se)[0]
    return g, r, b, se, ok


depths = [1, 2, 3, 4, 5, 6, 8, 10, 12, 15, 20, None]
xd = np.array([d if d else 30 for d in depths], dtype=float)
ax = axes[0]
g, r, b, se, ok = curves(ax, depths, 'max_depth', xd, 'max_depth')
simple = ok[np.argmin(xd[ok])]
ax.scatter([xd[simple]], [r['mean_test_score'][simple]], s=160, facecolors='none', edgecolors=GREEN, lw=1.8, zorder=7,
           label='simplest within 1 SE')
ax.set_xticks([1, 3, 5, 8, 10, 12, 15, 20, 30], ['1', '3', '5', '8', '10', '12', '15', '20', 'None'])
ax.set_xlabel('max_depth  (None = grow until the leaves are pure)')
ax.set_ylabel('accuracy')
ax.set_title(f'Pre-pruning: best depth {depths[b]}, CV {r["mean_test_score"][b]:.3f}', fontsize=10)
ax.legend(fontsize=7.5, loc='lower right', frameon=False)

alphas = np.r_[0, np.geomspace(1e-4, 1e-2, 21)]
leaves = np.array([DecisionTreeClassifier(random_state=0, ccp_alpha=a).fit(X_dev, y_dev).get_n_leaves()
                   for a in alphas])
ax = axes[1]
g, r, b, se, ok = curves(ax, alphas, 'ccp_alpha', leaves, 'leaves')
simple = ok[np.argmin(leaves[ok])]
ax.scatter([leaves[simple]], [r['mean_test_score'][simple]], s=160, facecolors='none', edgecolors=GREEN,
           lw=1.8, zorder=7)
ax.set_xscale('log')
ax.set_xticks([5, 10, 20, 50, 100, 200], ['5', '10', '20', '50', '100', '200'])
ax.set_xlabel('leaves of the pruned tree  (fewer = larger ccp_alpha)')
ax.annotate(f'ccp_alpha = 0\n{leaves[0]} leaves, CV {r["mean_test_score"][0]:.3f}', xy=(leaves[0], r['mean_test_score'][0]),
            xytext=(60, 0.80), fontsize=8, arrowprops=dict(arrowstyle='->', color='gray'))
ax.annotate(f'1-SE choice: ccp_alpha = {alphas[simple]:.5f}\n{leaves[simple]} leaves, CV {r["mean_test_score"][simple]:.3f}',
            xy=(leaves[simple], r['mean_test_score'][simple]), xytext=(5.6, 0.962), fontsize=8,
            arrowprops=dict(arrowstyle='->', color=GREEN))
ax.set_title(f'Post-pruning: best {leaves[b]} leaves, CV {r["mean_test_score"][b]:.3f}', fontsize=10)
ax.set_ylim(0.76, 1.01)
for a in axes:
    a.spines[['top', 'right']].set_visible(False)
plt.tight_layout()

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

plt.close()