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

# The lesson's gradient boosting model on California housing
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)
model = GradientBoostingRegressor(n_estimators=200, max_depth=4, random_state=0).fit(X_train, y_train)

# ICE by hand: set income to v in every test row, predict, keep every row's curve
grid = np.linspace(1, 10, 19)
ice = np.column_stack([model.predict(X_test.assign(income=v)) for v in grid])   # rows x grid
pdp = ice.mean(axis=0)                                                           # the PDP is their mean

old = (X_test['house_age'] >= 40).to_numpy()
new = (X_test['house_age'] < 20).to_numpy()
i2, i8 = np.argmin(abs(grid - 2)), np.argmin(abs(grid - 8))
gain = ice[:, i8] - ice[:, i2]
rng = np.random.default_rng(0)
show = rng.choice(len(X_test), 80, replace=False)                                # draw 80 of the curves

fig, (left, right) = plt.subplots(1, 2, figsize=(9.5, 3.8))
for i in show:
    c = '#f0883e' if old[i] else '#58a6ff' if new[i] else '#8b949e'
    left.plot(grid, ice[i], color=c, lw=0.6, alpha=0.55)
    right.plot(grid, ice[i] - ice[i, 0], color=c, lw=0.6, alpha=0.55)
for ax, curve in [(left, pdp), (right, pdp - pdp[0])]:
    ax.plot(grid, curve, color='#3fb950', lw=3, label='PDP = average of all 4,087 curves')
    ax.set_xlabel('income, set by hand to this value (\\$10,000)', fontsize=8.5)
    ax.tick_params(labelsize=8)
    ax.grid(alpha=0.25)
left.set_ylabel('predicted house value (\\$1,000)', fontsize=8.5)
left.set_title('ICE: one curve per district', fontsize=9.5, loc='left')
right.set_ylabel('change since income = 1 (\\$1,000)', fontsize=8.5)
right.set_title('Centered ICE: the same curves, all starting at 0', fontsize=9.5, loc='left')
right.plot([], [], color='#f0883e', lw=1.5,
           label=f'houses ≥ 40 years old: gain 2→8 = {gain[old].mean():.0f}')
right.plot([], [], color='#58a6ff', lw=1.5,
           label=f'houses < 20 years old: gain 2→8 = {gain[new].mean():.0f}')
right.legend(fontsize=7.5, loc='upper left', frameon=False)
left.legend(fontsize=7.5, loc='upper left', frameon=False)
fig.suptitle(f'Income from \\$20k to \\$80k: +{gain.mean():.0f}k on average (PDP), '
             f'+{np.percentile(gain, 5):.0f}k to +{np.percentile(gain, 95):.0f}k district by district '
             f'(5th–95th percentile)', fontsize=9.5)

plt.tight_layout()

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

plt.close()