import matplotlib.pyplot as plt
import numpy as np
from io import StringIO
from matplotlib.colors import LinearSegmentedColormap
from sklearn.datasets import make_blobs
from sklearn.linear_model import LogisticRegression

# The lesson's two-feature example: a straight boundary, parallel probability contours
X2, y2 = make_blobs(n_samples=200, centers=[[-1, -0.5], [1, 1]], cluster_std=1.0, random_state=0)
clf = LogisticRegression(C=np.inf).fit(X2, y2)
w, b = clf.coef_[0], clf.intercept_[0]
norm = np.linalg.norm(w)

fig, (ax, ax2) = plt.subplots(1, 2, figsize=(10, 4.2), gridspec_kw={'width_ratios': [1.1, 1]})

g1, g2 = np.meshgrid(np.linspace(-4, 4, 200), np.linspace(-3.5, 4, 200))
P = clf.predict_proba(np.c_[g1.ravel(), g2.ravel()])[:, 1].reshape(g1.shape)
ax.contourf(g1, g2, P, levels=np.linspace(0, 1, 11), cmap=LinearSegmentedColormap.from_list('gb', ['#8b949e', '#58a6ff']), alpha=0.2,
            rasterized=True)
levels = [0.1, 0.3, 0.7, 0.9]
cs = ax.contour(g1, g2, P, levels=levels, colors='gray', linewidths=0.8, linestyles='--')
# put the labels in the empty top-left corner, where each contour crosses x2 = 3.2
spots = [((np.log(p / (1 - p)) - b - w[1] * 3.2) / w[0], 3.2) for p in levels]
ax.clabel(cs, fontsize=7.5, fmt='%.1f', manual=spots)
ax.contour(g1, g2, P, levels=[0.5], colors='#f0883e', linewidths=2)
ax.scatter(*X2[y2 == 0].T, s=14, color='gray', alpha=0.8, label='class 0')
ax.scatter(*X2[y2 == 1].T, s=14, color='#58a6ff', alpha=0.9, label='class 1')
# the direction of w, drawn from a point of the boundary
p0 = -b * w / norm**2
ax.annotate('', xy=p0 + 1.5 * w / norm, xytext=p0,
            arrowprops=dict(arrowstyle='->', color='#3fb950', lw=1.8))
ax.text(*(p0 + 1.6 * w / norm + [0.05, 0.1]), 'w', color='#3fb950', fontsize=11, fontweight='bold')
ax.set_xlim(-4, 4)
ax.set_ylim(-3.5, 4)
ax.set_xlabel('x₁')
ax.set_ylabel('x₂')
ax.legend(fontsize=7.5, loc='lower right')
ax.set_title(f'boundary (orange): {w[0]:.2f}·x₁ + {w[1]:.2f}·x₂ − {-b:.2f} = 0\n'
             'probability contours 0.1 … 0.9 are parallel to it', fontsize=9.5, loc='left')
ax.tick_params(labelsize=8)

# right: probability depends only on the signed distance to the boundary
d = np.linspace(-3, 3, 300)
dist = (X2 @ w + b) / norm
ax2.plot(d, 1 / (1 + np.exp(-norm * d)), color='#58a6ff', lw=2, label=f'σ(|w|·d),  |w| = {norm:.2f}')
ax2.scatter(dist, clf.predict_proba(X2)[:, 1], s=10, color='gray', alpha=0.6,
            label='the 200 training points')
for dd in [0.5, 1, 2]:
    pp = 1 / (1 + np.exp(-norm * dd))
    ax2.plot([dd, dd], [0, pp], color='#3fb950', lw=0.8, ls=':')
    ax2.text(dd + 0.05, pp - 0.08, f'd = {dd:g}: {pp:.3f}', fontsize=7.5, color='#3fb950')
ax2.axvline(0, color='#f0883e', lw=1.5)
ax2.set_xlabel('signed distance to the boundary, d = (w·x + b) / |w|')
ax2.set_ylabel('predicted probability of class 1')
ax2.legend(fontsize=7.5, loc='upper left')
ax2.set_title('probability depends only on the distance to the line:\nthe further away, the more confident', fontsize=9.5, loc='left')
ax2.tick_params(labelsize=8)

plt.tight_layout()

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

plt.close()