import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from io import StringIO
from itertools import combinations
from math import factorial
from matplotlib.colors import LinearSegmentedColormap
from sklearn.ensemble import GradientBoostingRegressor
from sklearn.model_selection import train_test_split

# Exact Shapley values of 300 test districts (brute force over the 256 coalitions, vectorized),
# for the lesson's gradient boosting model and its 100-row background
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)
background = X_train.sample(100, random_state=0).to_numpy()
sample = X_test.sample(300, random_state=0)
A = sample.to_numpy()
n, d, m = len(A), A.shape[1], len(background)

v = {}
for size in range(d + 1):
    for S in combinations(range(d), size):
        Z = np.tile(background, (n, 1))
        if S:
            Z[:, S] = np.repeat(A[:, S], m, axis=0)
        v[S] = model.predict(pd.DataFrame(Z, columns=X.columns)).reshape(n, m).mean(axis=1)
phi = np.zeros((n, d))
for S, vS in v.items():
    for j in range(d):
        if j not in S:
            S_j = tuple(sorted(S + (j,)))
            phi[:, j] += factorial(len(S)) * factorial(d - len(S) - 1) / factorial(d) * (v[S_j] - vS)

mean_abs = np.abs(phi).mean(axis=0)
order = np.argsort(mean_abs)                         # least important at the bottom
cmap = LinearSegmentedColormap.from_list('lowhigh', ['#58a6ff', '#c9a0dc', '#ff7b72'])
rng = np.random.default_rng(0)

fig, (left, right) = plt.subplots(1, 2, figsize=(10, 4.4), gridspec_kw={'width_ratios': [2.2, 1]},
                                  sharey=True)
for row, j in enumerate(order):
    vals = phi[:, j]
    ranks = sample.iloc[:, j].rank(pct=True).to_numpy()          # color = the feature's value, low to high
    counts, edges = np.histogram(vals, bins=40)
    dens = counts[np.clip(np.digitize(vals, edges) - 1, 0, 39)] / counts.max()
    jitter = rng.uniform(-1, 1, n) * 0.38 * np.sqrt(dens)       # wider where more dots pile up
    left.scatter(vals, row + jitter, c=cmap(ranks), s=7, lw=0, rasterized=True)
left.axvline(0, color='#8b949e', lw=0.8)
left.set_yticks(range(d), X.columns[order], fontsize=8.5)
left.set_xlabel('SHAP value (\\$1,000): push away from the base value', fontsize=8.5)
left.set_title('Beeswarm: one dot per district and feature', fontsize=9.5, loc='left')
sm = plt.cm.ScalarMappable(cmap=cmap)
cb = fig.colorbar(sm, ax=left, fraction=0.03, pad=0.01, ticks=[0, 1])
cb.ax.set_yticklabels(['low', 'high'], fontsize=7.5)
cb.set_label('feature value', fontsize=8)
cb.outline.set_visible(False)

right.barh(range(d), mean_abs[order], color='#58a6ff', height=0.6)
for row, j in enumerate(order):
    right.text(mean_abs[j] + 1, row, f'{mean_abs[j]:.1f}', va='center', fontsize=8)
right.set_xlabel('mean |SHAP value| (\\$1,000)', fontsize=8.5)
right.set_title('Global importance', fontsize=9.5, loc='left')
right.set_xlim(0, mean_abs.max() * 1.2)
for ax in (left, right):
    ax.tick_params(axis='x', labelsize=8)
    ax.grid(axis='x', alpha=0.25)
    ax.set_axisbelow(True)
fig.suptitle(f'SHAP values of {n} test districts: the global picture is the local explanations, stacked',
             fontsize=10)

plt.tight_layout()

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

plt.close()