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

# The lesson's soft-margin experiment: same 100 points, three values of C
X, y = make_blobs(n_samples=100, centers=2, cluster_std=1.6, random_state=7)
cv = StratifiedKFold(n_splits=5, shuffle=True, random_state=0)

pad = 1.0
x0, x1 = X[:, 0].min() - pad, X[:, 0].max() + pad
y0, y1 = X[:, 1].min() - pad, X[:, 1].max() + pad
xx, yy = np.meshgrid(np.linspace(x0, x1, 300), np.linspace(y0, y1, 300))

fig, axes = plt.subplots(1, 3, figsize=(10, 3.9), sharey=True)
for ax, C in zip(axes, [0.01, 1, 100]):
    svm = SVC(kernel='linear', C=C).fit(X, y)
    width = 2 / np.linalg.norm(svm.coef_)
    acc = cross_val_score(SVC(kernel='linear', C=C), X, y, cv=cv).mean()
    Z = svm.decision_function(np.c_[xx.ravel(), yy.ravel()]).reshape(xx.shape)
    ax.contourf(xx, yy, Z, levels=[-1, 1], colors=['gray'], alpha=0.15)
    ax.contour(xx, yy, Z, levels=[-1, 0, 1], colors='gray', linestyles=['--', '-', '--'],
               linewidths=[1, 1.8, 1])
    for cls, c in zip([0, 1], ['#f0883e', '#58a6ff']):
        ax.scatter(X[y == cls, 0], X[y == cls, 1], s=14, color=c, zorder=3)
    sv = svm.support_vectors_
    ax.scatter(sv[:, 0], sv[:, 1], s=70, facecolors='none', edgecolors='#3fb950', lw=1.2, zorder=4)
    ax.set_title(f'C = {C:g}\nstreet width {width:.2f} · {len(sv)} support vectors\n'
                 f'5-fold CV accuracy {acc:.2f}', fontsize=9)
    ax.set_xlim(x0, x1); ax.set_ylim(y0, y1)
    ax.tick_params(labelsize=8)
    ax.set_xlabel('$x_1$', fontsize=9)
axes[0].set_xlabel('$x_1$\n\ncheap violations: wide street, many points inside', fontsize=8)
axes[2].set_xlabel('$x_1$\n\nexpensive violations: narrow street, few points on it', fontsize=8)
axes[0].set_ylabel('$x_2$', fontsize=9)
fig.suptitle('Soft margin: C sets the price of a point inside the street (circled = support vectors)',
             fontsize=10)

plt.tight_layout()

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

plt.close()