import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from io import StringIO
from sklearn.model_selection import train_test_split
from sklearn.neighbors import NearestNeighbors

# Where do the rows of a partial-dependence plot for latitude actually sit?
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)

real = NearestNeighbors(n_neighbors=1).fit(X[['latitude', 'longitude']])   # every real district
lats = np.arange(33, 42)
off = []
for lat in lats:
    dist, _ = real.kneighbors(X_test[['latitude', 'longitude']].assign(latitude=lat))
    off.append((dist[:, 0] > 0.25).mean())
off = np.array(off)

fig, (left, right) = plt.subplots(1, 2, figsize=(9.5, 4.3), gridspec_kw={'width_ratios': [1.1, 1]})
left.scatter(X['longitude'], X['latitude'], s=1, color='#8b949e', alpha=0.35, lw=0, rasterized=True)
examples = {5403: ('a San Francisco\ndistrict', -2.3), 9151: ('a Fresno\ndistrict', 0.35)}
for k, (idx, (label, dx)) in enumerate(examples.items()):
    lon = X.loc[idx, 'longitude']
    moved = pd.DataFrame({'latitude': lats.astype(float), 'longitude': lon})
    dist, _ = real.kneighbors(moved)
    ok = dist[:, 0] <= 0.25
    left.plot([lon, lon], [lats[0], lats[-1]], color='#8b949e', lw=0.8, ls=':')
    left.scatter(moved['longitude'][ok], moved['latitude'][ok], s=28, color='#3fb950', zorder=3,
                 label='on the map (a real district within 0.25°)' if k == 0 else None)
    left.scatter(moved['longitude'][~ok], moved['latitude'][~ok], s=28, color='#ff7b72', marker='X', zorder=3,
                 label='off the map' if k == 0 else None)
    left.scatter([lon], [X.loc[idx, 'latitude']], s=90, facecolors='none', edgecolors='#58a6ff', lw=1.8,
                 zorder=4)
    left.annotate(label, xy=(lon, X.loc[idx, 'latitude']), xytext=(lon + dx, X.loc[idx, 'latitude'] + 0.5),
                  fontsize=8, color='#58a6ff')
left.set_xlabel('longitude', fontsize=8.5)
left.set_ylabel('latitude', fontsize=8.5)
left.set_title('Each dot column: one district, latitude set to 33, 34, …, 41', fontsize=9, loc='left')
left.legend(fontsize=7.5, loc='lower left', frameon=False)
left.set_aspect('equal', adjustable='box')

right.barh(lats, off * 100, color=['#ff7b72' if f > 0.5 else '#f0883e' for f in off], height=0.65)
for lat, f in zip(lats, off):
    right.text(f * 100 + 1.5, lat, f'{f:.0%}', va='center', fontsize=8)
right.set_yticks(lats)
right.set_xlim(0, 100)
right.set_xlabel('% of the 4,087 test rows that land off the map', fontsize=8.5)
right.set_ylabel('latitude set to', fontsize=8.5)
right.set_title(f'On average, {off.mean():.0%} of the PDP\'s rows are fictitious',
                fontsize=9, loc='left')
for ax in (left, right):
    ax.tick_params(labelsize=8)
    ax.grid(alpha=0.2)
    ax.set_axisbelow(True)
fig.suptitle('Changing latitude alone moves districts into the Pacific and the Nevada desert', fontsize=10)

plt.tight_layout()

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

plt.close()