import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from io import StringIO
from sklearn.metrics import pairwise_distances
from sklearn.preprocessing import OneHotEncoder, OrdinalEncoder

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

ports = pd.DataFrame({'Embarked': ['C', 'Q', 'S']})       # Cherbourg, Queenstown, Southampton
codes = OrdinalEncoder().fit_transform(ports)              # C=0, Q=1, S=2
onehot = OneHotEncoder(sparse_output=False).fit_transform(ports)
d_int, d_oh = pairwise_distances(codes), pairwise_distances(onehot)
colors = ['#58a6ff', '#3fb950', '#f0883e']
names = ['C (Cherbourg)', 'Q (Queenstown)', 'S (Southampton)']

fig, (a1, a2) = plt.subplots(1, 2, figsize=(9.5, 3.4), gridspec_kw={'width_ratios': [1.25, 1]})
# left: the integer code puts the ports on a line
for i in range(3):
    a1.scatter(codes[i, 0], 0, s=120, color=colors[i], zorder=3, edgecolor='white')
    a1.text(codes[i, 0], -0.22, f'{names[i]}\ncode {codes[i, 0]:.0f}', ha='center', va='top', fontsize=8.5,
            color=colors[i])
a1.plot([0, 2], [0, 0], color='gray', lw=1)
for (i, j, h) in [(0, 1, 0.22), (1, 2, 0.22), (0, 2, 0.55)]:
    a1.annotate('', xy=(codes[j, 0], h), xytext=(codes[i, 0], h),
                arrowprops=dict(arrowstyle='<->', color='gray', lw=1))
    a1.text((codes[i, 0] + codes[j, 0]) / 2, h + 0.05, f'distance {d_int[i, j]:.0f}', ha='center', fontsize=8.5)
a1.set_xlim(-0.6, 2.6); a1.set_ylim(-0.9, 0.85); a1.axis('off')
a1.set_title('integer code: Q sits "between" C and S,\nand C–S is "twice as far" as C–Q', fontsize=9)
# right: one-hot puts every port on its own axis; drawn as the triangle joining (1,0,0), (0,1,0), (0,0,1)
tri = np.array([[0, 0], [1, 0], [0.5, np.sqrt(3) / 2]])
for i in range(3):
    j = (i + 1) % 3
    a2.plot(*zip(tri[i], tri[j]), color='gray', lw=1)
    mid = (tri[i] + tri[j]) / 2
    off = np.array([0, -0.09]) if i == 0 else (mid - tri.mean(0)) * 0.35
    a2.text(*(mid + off), f'{d_oh[i, j]:.3f}', ha='center', va='center', fontsize=8.5)
vec = ['(1, 0, 0)', '(0, 1, 0)', '(0, 0, 1)']
for i, off in enumerate([(-0.02, -0.16), (0.02, -0.16), (0, 0.1)]):
    a2.scatter(*tri[i], s=120, color=colors[i], zorder=3, edgecolor='white')
    a2.text(tri[i, 0] + off[0], tri[i, 1] + off[1], f'{names[i]}\n{vec[i]}', ha='center', va='center',
            fontsize=8.5, color=colors[i])
a2.set_xlim(-0.35, 1.35); a2.set_ylim(-0.35, 1.1); a2.set_aspect('equal'); a2.axis('off')
a2.set_title('one-hot: every pair of ports\nat the same distance √2', fontsize=9)
plt.tight_layout()

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

plt.close()