import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from io import StringIO
from matplotlib.patches import Circle

plt.rcParams.update({'text.color': 'gray', 'axes.labelcolor': 'gray', 'axes.edgecolor': 'gray',
                     'xtick.color': 'gray', 'ytick.color': 'gray'})

# Three real Titanic passengers, drawn in raw units and after standardization
url = "https://raw.githubusercontent.com/hsandmann/biblio/refs/heads/main/ml/aula04/titanic/train.csv"
df = pd.read_csv(url).dropna(subset=['Age'])
mu, sd = df[['Age', 'Fare']].mean(), df[['Age', 'Fare']].std(ddof=0)
trio = df.set_index('PassengerId').loc[[1, 117, 152], ['Age', 'Fare']]
labels = {1: 'A  (22 y, £7.25)', 117: 'B  (70.5 y, £7.75)', 152: 'C  (22 y, £66.60)'}
colors = {1: '#ff7b72', 117: '#f0883e', 152: '#58a6ff'}

fig, axes = plt.subplots(1, 2, figsize=(9.5, 4.4))
for ax, scaled in zip(axes, [False, True]):
    pts = (df[['Age', 'Fare']] - mu) / sd if scaled else df[['Age', 'Fare']]
    t = (trio - mu) / sd if scaled else trio
    ax.scatter(pts['Age'], pts['Fare'], s=5, color='gray', alpha=0.35, lw=0)
    a = t.loc[1].to_numpy()
    d = {pid: np.sqrt(((t.loc[pid].to_numpy() - a) ** 2).sum()) for pid in [117, 152]}
    nearest = min(d, key=d.get)
    ax.add_patch(Circle(a, d[nearest], fill=False, ls='--', lw=1.2, color=colors[nearest]))
    for pid in [1, 117, 152]:
        p = t.loc[pid].to_numpy()
        ax.scatter(*p, s=60, color=colors[pid], zorder=3, edgecolor='white', lw=0.8)
        off = (6, -13) if pid == 1 else (6, 6)
        ax.annotate(labels[pid], p, xytext=off, textcoords='offset points', fontsize=8.5,
                    color=colors[pid], fontweight='bold')
    for pid in [117, 152]:
        ax.plot(*zip(a, t.loc[pid].to_numpy()), color=colors[pid], lw=1, alpha=0.8)
    ax.set_aspect('equal')
    if scaled:
        ax.set_xlim(-2.4, 3.8); ax.set_ylim(-1.2, 2.6)
        ax.set_xlabel('Age (standardized)', fontsize=9); ax.set_ylabel('Fare (standardized)', fontsize=9)
        title = f'standardized: d(A,B) = {d[117]:.2f}, d(A,C) = {d[152]:.2f}\nnearest to A is C'
    else:
        ax.set_xlim(-5, 95); ax.set_ylim(-10, 110)
        ax.set_xlabel('Age (years)', fontsize=9); ax.set_ylabel('Fare (£)', fontsize=9)
        title = f'raw units: d(A,B) = {d[117]:.1f}, d(A,C) = {d[152]:.1f}\nnearest to A is B'
    ax.set_title(title, fontsize=9.5, loc='left')
    ax.tick_params(labelsize=8)
fig.suptitle('Same three passengers, same two columns: the nearest neighbor depends on the units',
             fontsize=10)
plt.tight_layout()

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

plt.close()