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.ensemble import RandomForestClassifier
from sklearn.inspection import permutation_importance
from sklearn.model_selection import train_test_split

# The lesson's importance experiment on German credit, with two columns of pure noise added:
# a serial number (1,000 distinct values) and a coin flip (2 values)
credit = fetch_openml(data_id=31, as_frame=True)
Xc = credit.data.copy()
for col in Xc.select_dtypes('category'):
    Xc[col] = Xc[col].cat.codes
yc = (credit.target == 'bad').astype(int)
rng = np.random.default_rng(0)
Xc['random_id'] = rng.permutation(len(Xc))
Xc['coin'] = rng.integers(0, 2, len(Xc))
Xc_tr, Xc_te, yc_tr, yc_te = train_test_split(Xc, yc, test_size=0.3, stratify=yc, random_state=0)

forest = RandomForestClassifier(n_estimators=300, random_state=0, n_jobs=-1).fit(Xc_tr, yc_tr)
perm = permutation_importance(forest, Xc_te, yc_te, scoring='roc_auc', n_repeats=10,
                              random_state=0, n_jobs=-1)
imp = pd.DataFrame({'MDI': forest.feature_importances_, 'perm': perm.importances_mean,
                    'std': perm.importances_std, 'k': Xc.nunique()}, index=Xc.columns)
order = imp.sort_values('MDI', ascending=False)
show = pd.concat([order.head(12), order.loc[['coin']]]) if 'coin' not in order.head(12).index else order.head(12)
show = show.iloc[::-1]                                  # largest at the top of the chart
noise = show.index.isin(['random_id', 'coin'])
colors = np.where(noise, '#ff7b72', '#58a6ff')
labels = [f'{name}  ({k})' for name, k in zip(show.index, show['k'])]

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(9.5, 4.2), sharey=True)
ax1.barh(labels, show['MDI'], color=colors)
ax1.set_title('Impurity-based (MDI), training data', fontsize=9, loc='left')
ax1.set_xlabel('mean decrease in impurity (sums to 1)', fontsize=8)
ax2.barh(labels, show['perm'], xerr=show['std'], color=colors, error_kw=dict(ecolor='gray', lw=0.8))
ax2.axvline(0, color='gray', lw=0.8)
ax2.set_title('Permutation, 300 held-out loans', fontsize=9, loc='left')
ax2.set_xlabel('drop in ROC AUC when the column is shuffled', fontsize=8)
for ax in (ax1, ax2):
    ax.tick_params(labelsize=8)
rank_mdi = int(order.index.get_loc('random_id')) + 1
rank_perm = int(imp['perm'].rank(ascending=False)['random_id'])
fig.suptitle(f'German credit + two noise columns (red): the serial number is #{rank_mdi} by MDI, '
             f'#{rank_perm} of 22 by permutation.  In brackets: distinct values', fontsize=9)

plt.tight_layout()

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

plt.close()