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

# The lesson's 12 points on a line: the outer ones are +1, the inner ones -1
x = np.array([-3.0, -2.6, -2.1, -1.2, -0.7, -0.2, 0.3, 0.9, 1.3, 2.0, 2.4, 2.9])
y = np.where(np.abs(x) > 1.6, 1, -1)
Phi = np.c_[x, x ** 2]                                  # the feature map x -> (x, x^2)
svm = SVC(kernel='linear', C=1e6, tol=1e-8).fit(Phi, y)
(w1, w2), b = svm.coef_[0], svm.intercept_[0]
roots = np.sort(np.roots([w2, w1, b]))                  # where the line meets the parabola

colors = np.where(y == 1, '#58a6ff', '#f0883e')
fig, (left, right) = plt.subplots(1, 2, figsize=(9.5, 3.6), gridspec_kw={'width_ratios': [1, 1.15]})

left.axhline(0, color='gray', lw=0.8)
left.scatter(x, np.zeros_like(x), c=colors, s=45, zorder=3)
for r in roots:
    left.axvline(r, color='#3fb950', ls='--', lw=1.3)
left.text(0, 0.55, 'no single threshold separates\nthe classes: the best one\ngets 9 of 12 right',
          ha='center', fontsize=8, color='gray')
left.text(0, -0.55, f'the lifted line, brought back:\ntwo thresholds, at {roots[0]:.2f} and {roots[1]:.2f}',
          ha='center', fontsize=8, color='#3fb950')
left.set_ylim(-1, 1); left.set_yticks([])
left.set_xlim(-3.4, 3.4)
left.set_xlabel('x  (original feature)', fontsize=9)
left.set_title('1 feature: not linearly separable', fontsize=10)

xs = np.linspace(-3.4, 3.4, 200)
right.plot(xs, xs ** 2, color='gray', lw=0.8, alpha=0.6)
right.scatter(x, x ** 2, c=colors, s=45, zorder=3)
right.plot(xs, -(w1 * xs + b) / w2, color='#3fb950', lw=1.8, label='SVM line in the lifted space')
right.plot(xs, -(w1 * xs + b - 1) / w2, color='#3fb950', lw=0.8, ls='--')
right.plot(xs, -(w1 * xs + b + 1) / w2, color='#3fb950', lw=0.8, ls='--')
sv = Phi[svm.support_]
right.scatter(sv[:, 0], sv[:, 1], s=140, facecolors='none', edgecolors='#3fb950', lw=1.3, zorder=4)
right.set_xlim(-3.4, 3.4); right.set_ylim(-0.5, 10)
right.set_xlabel('x', fontsize=9); right.set_ylabel('x²  (new feature)', fontsize=9)
right.legend(fontsize=7.5, loc='upper center', framealpha=0.3)
right.set_title('Add x² as a feature: a straight line separates', fontsize=10)
for ax in (left, right):
    ax.tick_params(labelsize=8)

plt.tight_layout()

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

plt.close()