import matplotlib.pyplot as plt
import numpy as np
from io import StringIO
from sklearn.datasets import make_blobs
from sklearn.svm import SVC

# 40 separable points; many lines classify all of them correctly
X, y = make_blobs(n_samples=40, centers=2, cluster_std=1.0, random_state=6)
Y = np.where(y == 1, 1, -1)
svm = SVC(kernel='linear', C=1e6, tol=1e-8).fit(X, Y)
w, b = svm.coef_[0], svm.intercept_[0]


def half_margin(w, b):
    """Distance from the line w.x + b = 0 to the closest point (negative if a point is misclassified)."""
    return np.min(Y * (X @ w + b)) / np.linalg.norm(w)


# three other separating lines: tilt the SVM direction, then place the line
# somewhere inside the gap that this direction leaves between the two classes
def rotate(v, deg):
    t = np.radians(deg)
    return np.array([[np.cos(t), -np.sin(t)], [np.sin(t), np.cos(t)]]) @ v


candidates = []
for deg, frac in [(25, 0.3), (-30, 0.5), (8, 0.8)]:
    wc = rotate(w / np.linalg.norm(w), deg)
    proj = X @ wc
    lo, hi = proj[Y == -1].max(), proj[Y == 1].min()      # the gap along wc
    assert hi > lo, 'this direction does not separate the classes'
    candidates.append((wc, -(lo + frac * (hi - lo))))


x0, x1 = X[:, 0].min() - 1, X[:, 0].max() + 1
y0, y1 = X[:, 1].min() - 1, X[:, 1].max() + 1
xs = np.linspace(x0, x1, 2)


def line(ax, w, b, off=0.0, **kw):
    ax.plot(xs, -(w[0] * xs + b - off) / w[1], **kw)


fig, axes = plt.subplots(1, 2, figsize=(9, 5.4))
colors = ['#f0883e', '#58a6ff']
for ax in axes:
    for cls, c in zip([-1, 1], colors):
        ax.scatter(X[Y == cls, 0], X[Y == cls, 1], s=22, color=c, zorder=3,
                   label=f'class {cls:+d}')
    ax.set_xlim(x0, x1); ax.set_ylim(y0, y1)
    ax.set_aspect('equal')
    ax.tick_params(labelsize=8)
    ax.set_xlabel('$x_1$', fontsize=9); ax.set_ylabel('$x_2$', fontsize=9)

ax = axes[0]
for (wc, bc), ls, c in zip(candidates, ['-', '--', ':'], ['#3fb950', '#ff7b72', 'gray']):
    line(ax, wc, bc, color=c, ls=ls, lw=1.6,
         label=f'closest point at {half_margin(wc, bc):.2f}')
ax.legend(fontsize=7.5, loc='upper center', bbox_to_anchor=(0.5, -0.13), ncol=2, frameon=False)
ax.set_title('Three lines, all 100% correct on training:\nhow close does each pass to a point?', fontsize=9.5)

ax = axes[1]
line(ax, w, b, color='gray', lw=1.8)
line(ax, w, b, off=1, color='gray', lw=1, ls='--')
line(ax, w, b, off=-1, color='gray', lw=1, ls='--')
# the street between the dashed lines
xx = np.linspace(x0, x1, 50)
ax.fill_between(xx, -(w[0] * xx + b - 1) / w[1], -(w[0] * xx + b + 1) / w[1],
                color='gray', alpha=0.15, lw=0)
sv = svm.support_vectors_
ax.scatter(sv[:, 0], sv[:, 1], s=150, facecolors='none', edgecolors='#3fb950', lw=1.6,
           zorder=4, label=f'{len(sv)} support vectors')
ax.legend(fontsize=7.5, loc='upper center', bbox_to_anchor=(0.5, -0.13), ncol=2, frameon=False)
ax.set_title(f'Maximum margin: closest point at {half_margin(w, b):.2f},\n'
             f'street width 2/‖w‖ = {2 / np.linalg.norm(w):.2f}', fontsize=9.5)

plt.tight_layout()

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

plt.close()