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 sklearn.ensemble import GradientBoostingRegressor
from sklearn.model_selection import train_test_split

# Exact (brute-force) Shapley values of two districts, for the lesson's gradient boosting
# model; the same numbers shap.TreeExplainer(model, data=background) returns
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)


def shapley(predict, x, background):
    names = list(x.index)
    d = len(names)
    v = {}
    for size in range(d + 1):
        for S in combinations(names, size):
            v[S] = predict(background.assign(**{j: x[j] for j in S})).mean()
    phi = pd.Series(0.0, index=names)
    for S, vS in v.items():
        for j in names:
            if j not in S:
                S_j = tuple(k for k in names if k in S or k == j)
                phi[j] += factorial(len(S)) * factorial(d - len(S) - 1) / factorial(d) * (v[S_j] - vS)
    return phi, v[()]


fmt = {'income': '{:.2f}', 'house_age': '{:.0f}', 'rooms': '{:.1f}', 'bedrooms': '{:.2f}',
       'population': '{:.0f}', 'occupancy': '{:.2f}', 'latitude': '{:.2f}', 'longitude': '{:.2f}'}
cases = [(5403, 'San Francisco district'), (9151, 'Fresno district')]
fig, axes = plt.subplots(1, 2, figsize=(10, 4.2))
for ax, (idx, title) in zip(axes, cases):
    x = X_test.loc[idx]
    phi, base = shapley(model.predict, x, background)
    pred = model.predict(x.to_frame().T)[0]
    order = list(phi.abs().sort_values(ascending=False).index)     # largest contribution on top
    start = base
    for k, name in enumerate(order):
        val = phi[name]
        color = '#ff7b72' if val > 0 else '#58a6ff'
        row = len(order) - k
        ax.barh(row, val, left=start, color=color, height=0.62)
        ax.text(max(start, start + val) + 4, row, f'{val:+.1f}', va='center', fontsize=8, color=color)
        start += val
    ax.axvline(base, color='#8b949e', lw=0.9, ls='--')
    ax.axvline(pred, color='#3fb950', lw=1.4)
    ax.set_yticks(range(len(order), 0, -1),
                  [f"{n} = {fmt[n].format(np.round(x[n], 2))}" for n in order], fontsize=8)
    ax.set_xlabel('predicted house value ($1,000)'.replace('$', r'\$'), fontsize=8.5)
    ax.set_title(f'{title}: {base:.0f} (base) → {pred:.0f} (prediction); actual {y_test[idx]:.0f}',
                 fontsize=9, loc='left')
    path = base + np.concatenate([[0], np.cumsum(phi[order].to_numpy())])
    span = path.max() - path.min()
    ax.set_xlim(path.min() - 0.08 * span, path.max() + 0.16 * span)
    ax.tick_params(axis='x', labelsize=8)
    ax.grid(axis='x', alpha=0.25)
    ax.set_axisbelow(True)
    ax.text(base, len(order) + 0.75, 'base value', fontsize=7.5, color='#8b949e', ha='center')
    ax.text(pred, len(order) + 0.75, 'prediction', fontsize=7.5, color='#3fb950', ha='center')
    ax.set_ylim(0.4, len(order) + 1.1)
fig.suptitle('Local explanations: each bar moves the prediction from the average district to this one '
             '(red up, blue down)', fontsize=10)

plt.tight_layout()

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

plt.close()