import matplotlib.pyplot as plt
import numpy as np
from io import StringIO
from sklearn.cluster import DBSCAN, HDBSCAN
from sklearn.datasets import make_blobs
from sklearn.metrics import adjusted_rand_score
from sklearn.neighbors import NearestNeighbors

# Two tight groups (std 0.2) side by side and one loose group (std 1.0).
# Left: the k-distance plot, the distance from every point to its 5th nearest
# neighbor (itself included, as DBSCAN counts with min_samples=5), sorted.
# Then DBSCAN with a small and a large eps, and HDBSCAN.
X, y = make_blobs(n_samples=[150, 150, 200], centers=[[0, 0], [1.2, 0], [4, 2]],
                  cluster_std=[0.2, 0.2, 1.0], random_state=0)
kdist = NearestNeighbors(n_neighbors=5).fit(X).kneighbors(X)[0][:, -1]
order = np.argsort(kdist)

BLUE, ORANGE, GREEN, RED = '#58a6ff', '#f0883e', '#3fb950', '#ff7b72'
PAL = np.array([BLUE, ORANGE, GREEN, '#d2a8ff', '#e3b341', '#56d4dd', '#ff9bce', '#a5d6ff', '#7ee787'])
plt.rcParams.update({'text.color': 'gray', 'axes.labelcolor': 'gray', 'axes.edgecolor': 'gray',
                     'xtick.color': 'gray', 'ytick.color': 'gray'})
fig, axes = plt.subplots(1, 4, figsize=(10.5, 3.1), gridspec_kw={'width_ratios': [1.25, 1, 1, 1]})
ax = axes[0]
group_col = np.array([BLUE, BLUE, GREEN])[y[order]]
ax.scatter(np.arange(len(X)), kdist[order], s=4, c=group_col, rasterized=True)
for eps, c in [(0.15, ORANGE), (0.4, RED)]:
    ax.axhline(eps, color=c, ls='--', lw=1)
    ax.text(5, eps + 0.03, f'eps = {eps}', color=c, fontsize=8)
ax.text(20, 0.62, 'blue: points of the two\ntight groups\ngreen: the loose group', fontsize=7.5)
ax.set_xlabel('points, sorted', fontsize=8.5)
ax.set_ylabel('distance to 5th neighbor', fontsize=8.5)
ax.set_title('k-distance: tight groups near 0.1,\nthe loose one near 0.3', fontsize=9)
ax.tick_params(labelsize=7.5)

models = [(DBSCAN(eps=0.15, min_samples=5), 'DBSCAN eps = 0.15'),
          (DBSCAN(eps=0.4, min_samples=5), 'DBSCAN eps = 0.4'),
          (HDBSCAN(min_cluster_size=15), 'HDBSCAN, no eps')]
for ax, (model, title) in zip(axes[1:], models):
    lab = model.fit_predict(X)
    noise = lab == -1
    ax.scatter(X[noise, 0], X[noise, 1], s=5, color='gray', alpha=0.5, rasterized=True)
    ax.scatter(X[~noise, 0], X[~noise, 1], s=5, c=PAL[lab[~noise] % len(PAL)], rasterized=True)
    ax.set_title(f'{title}\n{lab.max() + 1} clusters · {noise.sum()} noise · ARI {adjusted_rand_score(y, lab):.2f}',
                 fontsize=9)
    ax.set_xticks([])
    ax.set_yticks([])
plt.tight_layout()

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

plt.close()