import matplotlib.pyplot as plt
import numpy as np
from io import StringIO
from sklearn.linear_model import LogisticRegression

# The hours-of-study model of the lesson; three new students placed on its sigmoid
hours = np.arange(1, 12.5, 0.5)
passed = np.array([0] * 11 + [1, 0] + [1] * 10)
model = LogisticRegression(C=np.inf).fit(hours.reshape(-1, 1), passed)
w, b = model.coef_[0, 0], model.intercept_[0]

sigmoid = lambda z: 1 / (1 + np.exp(-z))
z = np.linspace(-8, 8, 400)

fig, ax = plt.subplots(figsize=(7.5, 3.4))
ax.axhspan(0.5, 1.05, xmin=0.5, color='#3fb950', alpha=0.07, lw=0)
ax.axhspan(-0.05, 0.5, xmax=0.5, color='gray', alpha=0.07, lw=0)
ax.axhline(0.5, color='gray', ls=':', lw=0.8)
ax.axvline(0, color='gray', ls=':', lw=0.8)
ax.plot(z, sigmoid(z), color='#58a6ff', lw=2)
for h, tx, ty in [(4.0, -6.9, 0.2), (6.0, 0.8, 0.1), (7.5, -3.8, 0.82)]:
    zi = w * h + b
    ax.scatter([zi], [sigmoid(zi)], s=40, color='#f0883e', zorder=5)
    ax.annotate(f'{h:g} h:  z = {zi:.2f}\nσ(z) = {sigmoid(zi):.3f}', xy=(zi, sigmoid(zi)),
                xytext=(tx, ty), fontsize=8, color='#f0883e',
                arrowprops=dict(arrowstyle='->', color='gray', lw=0.8))
ax.text(2.9, 0.55, 'z > 0  →  p > 0.5  →  class 1', fontsize=8, color='#3fb950')
ax.text(-7.8, 0.42, 'z < 0  →  p < 0.5  →  class 0', fontsize=8, color='gray')
ax.set_xlim(-8, 8)
ax.set_ylim(-0.05, 1.05)
ax.set_xlabel('score  z = w·x + b')
ax.set_ylabel('σ(z) = probability of class 1')
ax.tick_params(labelsize=8)
top = ax.secondary_xaxis('top', functions=(lambda zz: (zz - b) / w, lambda hh: w * hh + b))
top.set_xlabel(f'hours of study  (w = {w:.3f}, b = {b:.3f})', fontsize=8.5)
top.tick_params(labelsize=8)
ax.set_title('The sigmoid turns any score into a probability', fontsize=10, loc='left')

plt.tight_layout()

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

plt.close()