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

# California housing (the instructor's copy): predict the median house value, in $1000s.
# One fixed train/validation split; every model adds 1,000 trees WITHOUT early stopping,
# and we record the validation error after each tree.
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_tr, X_va, y_tr, y_va = train_test_split(X, y, test_size=0.25, random_state=0)


def curves(**params):
    m = HistGradientBoostingRegressor(max_iter=1000, early_stopping=False, random_state=0, **params)
    m.fit(X_tr, y_tr)
    rmse = lambda preds, t: np.array([np.sqrt(np.mean((p - t) ** 2)) for p in preds])
    return rmse(m.staged_predict(X_va), y_va), rmse(m.staged_predict(X_tr), y_tr)


it = np.arange(1, 1001)
fig, (left, right) = plt.subplots(1, 2, figsize=(10, 3.7))
colors = {1.0: '#ff7b72', 0.3: '#f0883e', 0.1: '#58a6ff', 0.03: '#3fb950'}
for eta, c in colors.items():
    va, tr = curves(learning_rate=eta)
    b = va.argmin()
    left.plot(it, va, color=c, lw=1.8, label=f'η = {eta:g}: best {va[b]:.1f} at {b + 1} trees')
    left.scatter([b + 1], [va[b]], color=c, s=30, zorder=4, edgecolor='white')
    if eta == 1.0:
        left.plot(it, tr, '--', color=c, lw=1.2, label=f'η = 1, training error ({tr[-1]:.1f} at 1,000)')
left.set_xscale('log')
left.set_ylim(0, 80)
left.set_xlabel('number of trees (log scale)', fontsize=8.5)
left.set_ylabel('validation RMSE ($1000s)', fontsize=8.5)
left.set_title('Learning rate: small steps need more trees, and go further', fontsize=9)
left.legend(fontsize=7.5, loc='lower left')

sizes = {2: '#ff7b72', 8: '#f0883e', 31: '#58a6ff', 127: '#3fb950'}
for leaves, c in sizes.items():
    va, _ = curves(learning_rate=0.1, max_leaf_nodes=leaves)
    b = va.argmin()
    right.plot(it, va, color=c, lw=1.8, label=f'{leaves} leaves: best {va[b]:.1f} at {b + 1} trees')
    right.scatter([b + 1], [va[b]], color=c, s=30, zorder=4, edgecolor='white')
right.set_xscale('log')
right.set_ylim(40, 80)
right.set_xlabel('number of trees (log scale)', fontsize=8.5)
right.set_title('Tree size (η = 0.1): bigger trees learn faster, and overfit sooner', fontsize=9)
right.legend(fontsize=7.5, loc='upper right')
for ax in (left, right):
    ax.tick_params(labelsize=7.5)
    ax.grid(alpha=0.25)

plt.tight_layout()

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

plt.close()