import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from io import StringIO
from sklearn.compose import ColumnTransformer
from sklearn.impute import SimpleImputer
from sklearn.metrics import mean_absolute_error
from sklearn.model_selection import train_test_split
from sklearn.neighbors import KNeighborsRegressor
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler, TargetEncoder

# 20 different train/test splits of California housing. For each split, two steps are fit
# either on ALL rows before the split (leak) or inside the pipeline (honest); the model is
# always trained on the training rows and scored on the test rows.
url = "https://raw.githubusercontent.com/hsandmann/biblio/refs/heads/main/ml/aula02/housing.csv"
housing = pd.read_csv(url)
X = housing.drop(columns='median_house_value')
y = housing['median_house_value']
num = ['longitude', 'latitude', 'housing_median_age', 'total_rooms',
       'total_bedrooms', 'population', 'households', 'median_income']
zone = housing['latitude'].round(2).astype(str) + '_' + housing['longitude'].round(2).astype(str)
Z = X[['median_income', 'housing_median_age']].assign(zone=zone)


def knn():
    return KNeighborsRegressor(n_neighbors=10)


res = {k: [] for k in ['sc_leak', 'sc_honest', 'te_leak', 'te_honest']}
all_scaled = make_pipeline(SimpleImputer(strategy='median'), StandardScaler()).fit_transform(X[num])
enc = TargetEncoder(target_type='continuous', random_state=0).fit(Z[['zone']], y)
Z_enc = Z.drop(columns='zone').assign(zone_price=enc.transform(Z[['zone']])[:, 0])
for seed in range(20):
    tr, te = train_test_split(np.arange(len(y)), test_size=0.2, random_state=seed)
    m = knn().fit(all_scaled[tr], y.iloc[tr])
    res['sc_leak'].append(mean_absolute_error(y.iloc[te], m.predict(all_scaled[te])))
    p = make_pipeline(SimpleImputer(strategy='median'), StandardScaler(), knn()).fit(X[num].iloc[tr], y.iloc[tr])
    res['sc_honest'].append(mean_absolute_error(y.iloc[te], p.predict(X[num].iloc[te])))
    m = make_pipeline(StandardScaler(), knn()).fit(Z_enc.iloc[tr], y.iloc[tr])
    res['te_leak'].append(mean_absolute_error(y.iloc[te], m.predict(Z_enc.iloc[te])))
    p = make_pipeline(
        ColumnTransformer([('zone', TargetEncoder(target_type='continuous', random_state=0), ['zone'])],
                          remainder='passthrough'),
        StandardScaler(), knn()).fit(Z.iloc[tr], y.iloc[tr])
    res['te_honest'].append(mean_absolute_error(y.iloc[te], p.predict(Z.iloc[te])))
res = {k: np.array(v) / 1000 for k, v in res.items()}

BLUE, ORANGE = '#58a6ff', '#f0883e'
plt.rcParams.update({'text.color': 'gray', 'axes.labelcolor': 'gray', 'axes.edgecolor': 'gray',
                     'xtick.color': 'gray', 'ytick.color': 'gray'})
fig, axes = plt.subplots(1, 2, figsize=(9, 3.9), sharey=True)
panels = [('sc', 'StandardScaler + imputer: never look at y'),
          ('te', 'TargetEncoder of the zone: built from y')]
for ax, (key, title) in zip(axes, panels):
    a, b = res[key + '_leak'], res[key + '_honest']
    for u, v in zip(a, b):
        ax.plot([0, 1], [u, v], color='gray', lw=0.6, alpha=0.5)
    ax.scatter(np.zeros(20), a, color=ORANGE, s=18, zorder=3)
    ax.scatter(np.ones(20), b, color=BLUE, s=18, zorder=3)
    ax.text(-0.12, a.mean(), f'mean\n\\${a.mean():.1f}k', ha='right', va='center', fontsize=8, color=ORANGE)
    ax.text(1.12, b.mean(), f'mean\n\\${b.mean():.1f}k', ha='left', va='center', fontsize=8, color=BLUE)
    ax.set_xticks([0, 1])
    ax.set_xticklabels(['fit on all rows\n(before the split)', 'fit inside the pipeline\n(train rows only)'],
                       fontsize=8)
    ax.set_xlim(-0.6, 1.6)
    ax.set_title(title, fontsize=9)
    ax.tick_params(labelsize=8)
    for s in ('top', 'right'):
        ax.spines[s].set_visible(False)
axes[0].set_ylabel('test MAE (thousands of dollars)', fontsize=8.5)
axes[0].set_ylim(0, 1.15 * max(v.max() for v in res.values()))
fig.suptitle('The same leak, 20 different splits: harmless for a scaler, huge for a target encoder',
             fontsize=10)

plt.tight_layout()

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

plt.close()