import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from io import StringIO
from sklearn.ensemble import HistGradientBoostingRegressor, RandomForestRegressor
from sklearn.model_selection import KFold, cross_val_score

# California housing, the same 5 folds for every configuration. 16 random configurations
# per family, plus the defaults. Forest: max_features and min_samples_leaf (100 trees).
# Boosting: learning rate 0.01-1, 50-1000 trees (no early stopping), 4-127 leaves, leaf size.
url = 'https://raw.githubusercontent.com/hsandmann/biblio/refs/heads/main/ml/aula02/housing.csv'
df = pd.read_csv(url)
y = df.pop('median_house_value') / 1000
X = df.astype({'ocean_proximity': 'category'})
X_codes = X.assign(ocean_proximity=X['ocean_proximity'].cat.codes)     # the forest needs numbers
cv = KFold(n_splits=5, shuffle=True, random_state=0)
rng = np.random.default_rng(0)


def score(model, data):
    return cross_val_score(model, data, y, cv=cv, scoring='r2').mean()


forest, boost = [], []
for _ in range(16):
    rf = RandomForestRegressor(n_estimators=100, n_jobs=-1, random_state=0,
                               max_features=rng.choice([0.2, 0.33, 0.5, 0.7, 1.0]),
                               min_samples_leaf=int(rng.choice([1, 2, 5, 10, 20])))
    forest.append(score(rf, X_codes))
    hgb = HistGradientBoostingRegressor(early_stopping=False, random_state=0,
                                        learning_rate=10 ** rng.uniform(-2, 0),
                                        max_iter=int(rng.choice([50, 100, 300, 1000])),
                                        max_leaf_nodes=int(rng.choice([4, 8, 16, 31, 63, 127])),
                                        min_samples_leaf=int(rng.choice([5, 20, 50, 100])))
    boost.append(score(hgb, X))
default_rf = score(RandomForestRegressor(n_jobs=-1, random_state=0), X_codes)
default_hgb = score(HistGradientBoostingRegressor(random_state=0), X)

fig, ax = plt.subplots(figsize=(8, 3))
for row, (name, vals, dflt, c) in enumerate([('random forest', forest, default_rf, '#58a6ff'),
                                             ('gradient boosting\n(HistGradientBoosting)', boost, default_hgb, '#3fb950')]):
    vals = np.array(vals)
    jitter = np.random.default_rng(row).uniform(-0.12, 0.12, len(vals))
    ax.scatter(vals, row + jitter, color=c, s=28, alpha=0.8, lw=0, label=None)
    ax.scatter([dflt], [row], marker='*', s=190, color=c, edgecolor='gray', zorder=4)
    ax.text(0.01, (row + 0.78) / 2.1,  f'16 random configurations: R² from {vals.min():.3f} to {vals.max():.3f}'
            f' · default ★ {dflt:.3f}', fontsize=7.8, color='gray', va='center',
            transform=ax.transAxes)
ax.set_yticks([0, 1], ['random forest', 'gradient boosting\n(HistGradientBoosting)'], fontsize=8)
ax.set_ylim(-0.5, 1.6)
ax.set_xlabel('5-fold CV R² on California housing (same folds for every point)', fontsize=8.5)
ax.tick_params(labelsize=7.5)
ax.grid(axis='x', alpha=0.25)
ax.set_title('The forest is hard to get wrong; boosting has a higher ceiling and a lower floor', fontsize=9.5)

plt.tight_layout()

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

plt.close()