import matplotlib.pyplot as plt
import numpy as np
from io import StringIO
from imblearn.over_sampling import SMOTE
from sklearn.neighbors import NearestNeighbors

# A 2-D toy: 150 majority points and 12 minority points
rng = np.random.default_rng(5)
X_maj = rng.normal([0, 0], 1.0, size=(150, 2))
X_min = rng.normal([2.4, 2.0], 0.9, size=(12, 2))
X = np.vstack([X_maj, X_min])
y = np.r_[np.zeros(150, int), np.ones(12, int)]

# SMOTE (imbalanced-learn) asked for 30 minority points: 18 synthetic ones are appended
X_res, y_res = SMOTE(sampling_strategy={1: 30}, k_neighbors=5, random_state=0).fit_resample(X, y)
synth = X_res[len(X):]

# Recover, for each synthetic point, the segment x_i -> x_nn it was drawn on
# (x_nn among the 5 nearest minority neighbors of x_i) and its lambda
nn = NearestNeighbors(n_neighbors=6).fit(X_min).kneighbors(X_min, return_distance=False)[:, 1:]
segments = []
for s in synth:
    best = None
    for i in range(len(X_min)):
        for j in nn[i]:
            d = X_min[j] - X_min[i]
            lam = np.clip((s - X_min[i]) @ d / (d @ d), 0, 1)
            err = np.linalg.norm(X_min[i] + lam * d - s)
            if best is None or err < best[0]:
                best = (err, i, j, lam)
    segments.append(best)
max_err = max(b[0] for b in segments)

fig, ax = plt.subplots(figsize=(7.6, 5.6))
ax.scatter(*X_maj.T, s=10, color='#58a6ff', alpha=0.55, label='majority (150)')
for _, i, j, _ in segments:
    ax.plot(*np.c_[X_min[i], X_min[j]], color='gray', lw=0.8, ls='--', alpha=0.7)
ax.scatter(*X_min.T, s=45, color='#ff7b72', edgecolor='white', lw=0.6, zorder=3, label='minority, original (12)')
ax.scatter(*synth.T, s=38, marker='X', color='#f0883e', zorder=4, label='minority, synthetic (18)')

# one example spelled out
err, i, j, lam = segments[0]
xi, xj, s = X_min[i], X_min[j], synth[0]
ax.annotate('$x_i$', xy=xi, xytext=xi + [-0.5, 0.35], fontsize=10, color='#ff7b72',
            arrowprops=dict(arrowstyle='->', color='#ff7b72', lw=0.8))
ax.annotate('$x_{nn}$', xy=xj, xytext=xj + [0.35, 0.35], fontsize=10, color='#ff7b72',
            arrowprops=dict(arrowstyle='->', color='#ff7b72', lw=0.8))
ax.annotate(f'$x_{{new}} = x_i + \\lambda\\,(x_{{nn}} - x_i)$\n$\\lambda$ = {lam:.2f}', xy=s,
            xytext=(-2.9, 3.6), fontsize=9, color='#f0883e',
            arrowprops=dict(arrowstyle='->', color='#f0883e', lw=0.8))
ax.text(-2.9, -3.3, 'each dashed segment joins a minority point to one of its 5 nearest\n'
        f'minority neighbors; every synthetic point lies on one (max. gap {max_err:.1g})',
        fontsize=8, color='gray')
ax.set_xlim(-3, 5.2)
ax.set_ylim(-3.5, 4.6)
ax.set_aspect('equal')
ax.set_xlabel('feature 1')
ax.set_ylabel('feature 2')
ax.legend(fontsize=8, loc='lower right', bbox_to_anchor=(1.0, 0.1))
ax.grid(color='gray', alpha=0.2, lw=0.5)
ax.set_title('SMOTE: new minority points are interpolated, not copied', fontsize=10)

plt.tight_layout()

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

plt.close()