import matplotlib.pyplot as plt
import numpy as np
from io import StringIO

plt.rcParams.update({'text.color': '#808080', 'axes.labelcolor': '#808080', 'axes.edgecolor': '#999999',
                     'xtick.color': '#808080', 'ytick.color': '#808080'})

# The six points of the lesson and the three updates of the perceptron (w = 0, b = 0, eta = 1)
names = list('ABCDEF')
X = np.array([[4, 2], [0, 2], [1, 3], [2, 4], [3, 0], [3, 1]])
y = np.array([+1, -1, -1, -1, +1, +1])

stages = [((0, 0), 0, None, 'start: w = (0, 0), b = 0\nno boundary yet'),
          ((4, 2), 1, 0, 'after A: w = (4, 2), b = 1\n4x₁ + 2x₂ + 1 = 0'),
          ((4, 0), 0, 1, 'after B: w = (4, 0), b = 0\nx₁ = 0'),
          ((3, -3), -1, 2, 'after C: w = (3, −3), b = −1\nx₂ = x₁ − 1/3 · all correct')]

lo, hi = -1.2, 5.2
xx, yy = np.meshgrid(np.linspace(lo, hi, 300), np.linspace(lo, hi, 300))
fig, axes = plt.subplots(1, 4, figsize=(10, 3.1))
for ax, (w, b, upd, title) in zip(axes, stages):
    w = np.array(w, float)
    if w.any():
        side = w[0] * xx + w[1] * yy + b > 0
        ax.contourf(xx, yy, side, levels=[-0.5, 0.5, 1.5], colors=['#f0883e', '#58a6ff'], alpha=0.12).set_rasterized(True)
        ax.contour(xx, yy, w[0] * xx + w[1] * yy + b, levels=[0], colors='gray', linewidths=1.6)
        # w is perpendicular to the boundary and points to the +1 side
        c = np.array([1.5, 1.5])                    # anchor the arrow at the boundary point closest to (1.5, 1.5)
        p0 = c - (w @ c + b) / (w @ w) * w
        d = w / np.linalg.norm(w) * 1.0
        if lo < p0[0] < hi and lo < p0[1] < hi:
            ax.annotate('', xy=p0 + d, xytext=p0, arrowprops=dict(arrowstyle='->', color='#3fb950', lw=1.6))
            ax.text(*(p0 + d * 1.25), 'w', color='#3fb950', fontsize=9, ha='center', va='center')
    wrong = y * (X @ w + b) <= 0
    for (a, c), t, n, bad in zip(X, y, names, wrong):
        ax.scatter(a, c, s=55, color='#58a6ff' if t > 0 else '#f0883e', zorder=3,
                   marker='o' if t > 0 else 's')
        ax.text(a + 0.22, c + 0.22, n, fontsize=8.5)
        if bad:
            ax.scatter(a, c, s=150, facecolors='none', edgecolors='#ff7b72', lw=1.3, zorder=4)
    if upd is not None:
        ax.scatter(*X[upd], s=320, facecolors='none', edgecolors='#3fb950', lw=1.8, zorder=4)
    ax.set_xlim(lo, hi); ax.set_ylim(lo, hi); ax.set_aspect('equal')
    ax.set_xticks([0, 2, 4]); ax.set_yticks([0, 2, 4])
    ax.tick_params(labelsize=7.5)
    ax.set_title(title, fontsize=8.5)
    ax.text(lo + 0.15, lo + 0.2, f'{wrong.sum()} wrong or on the line', fontsize=7.5, color='#ff7b72')
axes[0].set_ylabel('x₂', fontsize=9)
for ax in axes:
    ax.set_xlabel('x₁', fontsize=9)
fig.suptitle('Perceptron, η = 1: three updates (green ring = the point that caused it; red ring = wrong side '
             'or on the line)', fontsize=9)
plt.tight_layout()

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

plt.close()