import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from io import StringIO
from sklearn.ensemble import RandomForestRegressor
from sklearn.inspection import permutation_importance
from sklearn.model_selection import train_test_split

# California housing, the lesson's 8 features, plus two columns of pure noise:
# a continuous one (thousands of distinct values) and a coin flip (two values)
url = 'https://raw.githubusercontent.com/hsandmann/biblio/refs/heads/main/ml/aula02/housing.csv'
df = pd.read_csv(url).dropna()
X = pd.DataFrame({
    'income': df['median_income'], 'house_age': df['housing_median_age'],
    'rooms': df['total_rooms'] / df['households'], 'bedrooms': df['total_bedrooms'] / df['households'],
    'population': df['population'], 'occupancy': df['population'] / df['households'],
    'latitude': df['latitude'], 'longitude': df['longitude']})
y = df['median_house_value'] / 1000
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

rng = np.random.default_rng(0)
X_train_n = X_train.assign(noise=rng.normal(size=len(X_train)), coin=rng.integers(0, 2, len(X_train)))
X_test_n = X_test.assign(noise=rng.normal(size=len(X_test)), coin=rng.integers(0, 2, len(X_test)))
forest = RandomForestRegressor(n_estimators=100, random_state=0, n_jobs=-1).fit(X_train_n, y_train)
on_train = permutation_importance(forest, X_train_n, y_train, n_repeats=5, random_state=0)
on_test = permutation_importance(forest, X_test_n, y_test, n_repeats=5, random_state=0)

table = pd.DataFrame({'mdi': forest.feature_importances_, 'train': on_train.importances_mean,
                      'test': on_test.importances_mean}, index=X_train_n.columns)
table = table.sort_values('test')
names = table.index.tolist()
pos = np.arange(len(names))
noise_rows = [names.index('noise'), names.index('coin')]

fig, (left, right) = plt.subplots(1, 2, figsize=(9.5, 3.9), sharey=True)
colors = ['#f0883e' if n in ('noise', 'coin') else '#8b949e' for n in names]
left.barh(pos, table['mdi'], color=colors, height=0.6)
for i in noise_rows:
    left.text(table['mdi'].iloc[i] + 0.01, i, f"{table['mdi'].iloc[i]:.3f}", va='center',
              fontsize=8, color='#f0883e')
left.set_yticks(pos, names, fontsize=8.5)
left.set_xlabel('impurity importance (MDI, shares sum to 1)', fontsize=8.5)
left.set_title('MDI: computed on the training data', fontsize=9.5, loc='left')

h = 0.38
right.barh(pos + h / 2, table['train'], height=h, color='#f0883e', label='shuffled in the training set')
right.barh(pos - h / 2, table['test'], height=h, color='#58a6ff', label='shuffled in the test set')
for i in noise_rows:
    right.text(max(table['train'].iloc[i], 0) + 0.015, i + h / 2, f"{table['train'].iloc[i]:.3f}",
               va='center', fontsize=7.5, color='#f0883e')
    right.text(max(table['test'].iloc[i], 0) + 0.015, i - h / 2, f"{abs(table['test'].iloc[i]):.3f}",
               va='center', fontsize=7.5, color='#58a6ff')
right.set_xlabel('permutation importance (drop in R² when shuffled)', fontsize=8.5)
right.set_title('Permutation: train vs held-out test', fontsize=9.5, loc='left')
right.legend(fontsize=7.5, loc='lower right', frameon=False)
for ax in (left, right):
    ax.tick_params(labelsize=8)
    ax.grid(axis='x', alpha=0.25)
    ax.set_axisbelow(True)
fig.suptitle('A random forest on California housing, with two columns of pure noise added (orange names)',
             fontsize=10)
for lab in left.get_yticklabels():
    if lab.get_text() in ('noise', 'coin'):
        lab.set_color('#f0883e')

plt.tight_layout()

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

plt.close()