import matplotlib.pyplot as plt
import numpy as np
from io import StringIO
from sklearn.datasets import make_circles, make_moons
from sklearn.model_selection import StratifiedKFold, cross_val_score
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler
from sklearn.svm import SVC

# Same three kernels as the lesson's table, on two datasets no line can separate
datasets = {
    'circles': make_circles(n_samples=200, factor=0.4, noise=0.12, random_state=0),
    'moons': make_moons(n_samples=200, noise=0.25, random_state=0),
}
kernels = {
    'linear': SVC(kernel='linear'),
    'polynomial, degree 3': SVC(kernel='poly', degree=3, coef0=1),
    'RBF': SVC(kernel='rbf'),
}
cv = StratifiedKFold(n_splits=5, shuffle=True, random_state=0)

fig, axes = plt.subplots(2, 3, figsize=(9.5, 6.2))
for row, (dname, (X, y)) in enumerate(datasets.items()):
    pad = 0.4
    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, 160), np.linspace(y0, y1, 160))
    for col, (kname, svc) in enumerate(kernels.items()):
        ax = axes[row, col]
        model = make_pipeline(StandardScaler(), svc)
        acc = cross_val_score(model, X, y, cv=cv).mean()
        model.fit(X, y)
        Z = model.decision_function(np.c_[xx.ravel(), yy.ravel()]).reshape(xx.shape)
        ax.contourf(xx, yy, Z > 0, levels=[-0.5, 0.5, 1.5], colors=['#f0883e', '#58a6ff'], alpha=0.13).set_rasterized(True)
        ax.contour(xx, yy, Z, levels=[0], colors='gray', linewidths=1.4)
        ax.contour(xx, yy, Z, levels=[-1, 1], colors='gray', linewidths=0.7, linestyles='--')
        ax.scatter(X[y == 0, 0], X[y == 0, 1], s=9, color='#f0883e')
        ax.scatter(X[y == 1, 0], X[y == 1, 1], s=9, color='#58a6ff')
        ax.set_xlim(x0, x1); ax.set_ylim(y0, y1)
        ax.set_xticks([]); ax.set_yticks([])
        ax.set_title(f'{kname} — CV accuracy {acc:.3f}', fontsize=9)
        if col == 0:
            ax.set_ylabel(dname, fontsize=10)
fig.suptitle('Same SVM, three kernels (scaled inputs, C = 1, 5-fold CV)', fontsize=10)

plt.tight_layout()

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

plt.close()