import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from io import StringIO
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.preprocessing import StandardScaler

# The lesson's manual workflow on California housing, and the three classic bugs.
# The model and the training are identical in every bar: only the prediction code changes.
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']
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
num = ['longitude', 'latitude', 'housing_median_age', 'total_rooms',
       'total_bedrooms', 'population', 'households', 'median_income']

imputer = SimpleImputer(strategy='median')
scaler = StandardScaler()
model = KNeighborsRegressor(n_neighbors=10)
model.fit(scaler.fit_transform(imputer.fit_transform(X_train[num])), y_train)


def mae(p):
    return mean_absolute_error(y_test, p) / 1000


def predict_bug1(batch):
    return model.predict(StandardScaler().fit_transform(imputer.transform(batch[num])))


correct = mae(model.predict(scaler.transform(imputer.transform(X_test[num]))))
whole = mae(predict_bug1(X_test))
by_region = pd.concat([pd.Series(predict_bug1(b), index=b.index) for _, b in X_test.groupby('ocean_proximity')])
regions = mae(by_region[X_test.index])
one = mae([predict_bug1(X_test.iloc[[i]])[0] for i in range(len(X_test))])
no_scaler = mae(model.predict(imputer.transform(X_test[num])))
swapped = num[:5] + ['households', 'population', 'median_income']
swap = mae(model.predict(scaler.transform(imputer.transform(X_test[swapped].to_numpy()))))
baseline = mae(np.full(len(y_test), y_train.mean()))

BLUE, ORANGE, RED = '#58a6ff', '#f0883e', '#ff7b72'
plt.rcParams.update({'text.color': 'gray', 'axes.labelcolor': 'gray', 'axes.edgecolor': 'gray',
                     'xtick.color': 'gray', 'ytick.color': 'gray'})
rows = [
    ('careful manual workflow (fit on train, transform the rest)', correct, BLUE),
    ('bug 1 · fit_transform on the whole test set', whole, ORANGE),
    ("bug 1 · fit_transform on each region's batch", regions, ORANGE),
    ('bug 1 · fit_transform on one district at a time', one, RED),
    ('bug 2 · scaler forgotten at prediction time', no_scaler, RED),
    ('bug 3 · two columns swapped, sent as an array', swap, RED),
]
fig, ax = plt.subplots(figsize=(9, 3.9))
ypos = np.arange(len(rows))[::-1]
for yy, (label, v, c) in zip(ypos, rows):
    ax.barh(yy, v, color=c, height=0.62)
    ax.text(v + 1.5, yy, f'${v:,.1f}k', va='center', fontsize=8.5)
ax.axvline(baseline, color='gray', ls='--', lw=1.2)
ax.text(baseline - 1.5, ypos[0] + 0.5, f'always predict the training mean: ${baseline:,.1f}k',
        fontsize=8, va='center', ha='right')
ax.set_yticks(ypos)
ax.set_yticklabels([r[0] for r in rows], fontsize=8.5)
ax.set_xlim(0, 140)
ax.set_ylim(-0.6, len(rows) - 0.05)
ax.set_xlabel('mean absolute error on the 4,128 test districts (thousands of dollars)', fontsize=8.5)
ax.tick_params(labelsize=8)
for s in ('top', 'right'):
    ax.spines[s].set_visible(False)
ax.set_title('Same trained model, same test districts: only the prediction code changes', fontsize=10)

plt.tight_layout()

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

plt.close()