import matplotlib.pyplot as plt
import numpy as np
from io import StringIO
from matplotlib.patches import Ellipse

# Six customers and one query (R$ 5,000, 25 years). Left: raw units, where one real of
# income weighs as much as one year of age. Right: after StandardScaler.
plt.rcParams.update({'text.color': 'gray', 'axes.labelcolor': 'gray', 'axes.edgecolor': 'gray',
                     'xtick.color': 'gray', 'ytick.color': 'gray'})
BLUE, ORANGE, GREEN = '#58a6ff', '#f0883e', '#3fb950'

names = ['Rita', 'Otávio', 'Sônia', 'Caio', 'Lívia', 'Mário']
income = np.array([5200, 5700, 9000, 3000, 7400, 2600])
age = np.array([61, 26, 45, 33, 52, 22])
bought = np.array([0, 1, 0, 1, 0, 1])
q = np.array([5000, 25])

P = np.c_[income, age].astype(float)
mu, sd = P.mean(axis=0), P.std(axis=0)
Z, zq = (P - mu) / sd, (q - mu) / sd
d_raw = np.sqrt(((P - q) ** 2).sum(axis=1))
d_z = np.sqrt(((Z - zq) ** 2).sum(axis=1))

fig, (left, right) = plt.subplots(1, 2, figsize=(10, 4.4))
for ax, pts, qq, d, title in [
        (left, P, q, d_raw, 'Raw units: nearest is Rita (61 years old)'),
        (right, Z, zq, d_z, 'Standardized: nearest is Otávio (26 years old)')]:
    for cls, color, label in [(1, BLUE, 'bought'), (0, ORANGE, 'did not buy')]:
        m = bought == cls
        ax.scatter(pts[m, 0], pts[m, 1], s=60, color=color, label=label, zorder=3)
    for i, n in enumerate(names):
        ax.annotate(f'{n} ({d[i]:.0f})' if ax is left else f'{n} ({d[i]:.2f})', pts[i],
                    xytext=(6, 5), textcoords='offset points', fontsize=8)
    ax.scatter(*qq, marker='*', s=220, color=GREEN, zorder=4, label='new customer')
    j = np.argmin(d)
    ax.plot([qq[0], pts[j, 0]], [qq[1], pts[j, 1]], color=GREEN, lw=1.5, zorder=2)
    # every point at the same distance as the nearest neighbor
    ax.add_patch(Ellipse(qq, 2 * d[j], 2 * d[j], fill=False, ls='--', color='gray', lw=1.2))
    ax.set_title(title, fontsize=10)
    ax.tick_params(labelsize=8)

left.set_xlabel('income (R$ per month)')
left.set_ylabel('age (years)')
left.set_xlim(1500, 10500)
left.set_ylim(10, 75)
left.text(5400, 72, '← the circle of radius 203 is this band:\n    ±R$ 203 of income, any age', fontsize=8, ha='left', va='top')
right.set_xlabel('income (standardized)')
right.set_ylabel('age (standardized)')
right.set_xlim(-1.9, 2.4)
right.set_ylim(-1.7, 1.9)
right.set_aspect('equal')
right.legend(fontsize=8, loc='upper left', frameon=False)

plt.tight_layout()

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

plt.close()