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

# LIME along a single feature, for the San Francisco district of the lesson:
# perturb income, weight each perturbed point by its closeness, fit a weighted line
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)

x = X_test.loc[5403]
sd = X_train['income'].std(ddof=0)
rng = np.random.default_rng(0)
steps = rng.normal(size=3000)                                  # perturbations, in standard deviations
Z = pd.DataFrame(np.tile(x.to_numpy(), (len(steps), 1)), columns=X.columns)
Z['income'] = x['income'] + steps * sd
f_Z = model.predict(Z)

grid = np.sort(np.concatenate([np.linspace(0.5, 9, 1000), Z["income"]]))   # include the perturbed points
curve = model.predict(pd.DataFrame(np.tile(x.to_numpy(), (len(grid), 1)), columns=X.columns).assign(income=grid))
f_x = model.predict(x.to_frame().T)[0]

fig, ax = plt.subplots(figsize=(8.5, 4))
ax.plot(grid, curve, color='#8b949e', lw=2, label='the model, income varied, everything else fixed')
widths = [(0.25, '#3fb950'), (0.75, '#f0883e'), (3.0, '#58a6ff')]
for w, c in widths:
    weight = np.sqrt(np.exp(-steps ** 2 / w ** 2))
    fit = Ridge(alpha=1).fit(steps.reshape(-1, 1), f_Z, sample_weight=weight)
    slope = fit.coef_[0]                                         # $1,000 per standard deviation of income
    s = np.linspace(-1.2, 1.2, 2)
    ax.plot(x['income'] + s * sd, f_x + slope * s, color=c, lw=2.2,
            label=f'kernel width {w} sd: slope {slope:+.0f}k per sd ({slope / sd:+.0f}k per \\$10k)')
    if w == 0.25:
        keep = weight > 0.05
        ax.scatter(Z['income'][keep][:250], f_Z[keep][:250], s=40 * weight[keep][:250], color=c, alpha=0.5, rasterized=True,
                   lw=0, zorder=2)
ax.scatter([x['income']], [f_x], s=80, color='#ff7b72', zorder=5, label=f'the district: income {x["income"]:.2f}')
ax.set_xlabel('income (\\$10,000)', fontsize=9)
ax.set_ylabel('predicted house value (\\$1,000)', fontsize=9)
ax.set_xlim(0.5, 9)
ax.tick_params(labelsize=8)
ax.grid(alpha=0.25)
ax.legend(fontsize=7.5, loc='lower right', frameon=False)
ax.set_title('One-feature LIME: the answer is a slope, and the slope depends on how wide "local" is\n'
             '(green dots: perturbed points of the narrowest kernel, size = weight)', fontsize=9.5, loc='left')

plt.tight_layout()

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

plt.close()