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

# Hours of study -> pass (1) / fail (0). Same data as the lesson's first code block.
hours = np.arange(1, 12.5, 0.5)
passed = np.array([0] * 11 + [1, 0] + [1] * 10)
far = np.array([25, 30, 35, 40, 45])

cases = [('23 students', hours, passed, 13),
         ('the same 23 + 5 students who studied a lot (and passed)',
          np.r_[hours, far], np.r_[passed, [1] * 5], 47)]

fig, axes = plt.subplots(1, 2, figsize=(10, 3.6), gridspec_kw={'width_ratios': [1, 1.6]})
for ax, (title, h, y, xmax) in zip(axes, cases):
    X = h.reshape(-1, 1)
    line = LinearRegression().fit(X, y)
    logit = LogisticRegression(C=np.inf).fit(X, y)
    xs = np.linspace(0, xmax, 400).reshape(-1, 1)
    cut_line = (0.5 - line.intercept_) / line.coef_[0]
    cut_log = -logit.intercept_[0] / logit.coef_[0, 0]
    wrong = (line.predict(X) >= 0.5) != y

    ax.axhspan(1, 1.6, color='#ff7b72', alpha=0.08, lw=0)
    ax.axhspan(-0.6, 0, color='#ff7b72', alpha=0.08, lw=0)
    ax.axhline(0.5, color='gray', ls=':', lw=0.8)
    ax.plot(xs, line.predict(xs), color='#f0883e', lw=1.8, label='linear regression')
    ax.plot(xs, logit.predict_proba(xs)[:, 1], color='#58a6ff', lw=1.8, label='logistic regression')
    ax.scatter(h[y == 1], y[y == 1], s=22, color='#3fb950', zorder=4, label='passed (y = 1)')
    ax.scatter(h[y == 0], y[y == 0], s=22, color='gray', zorder=4, label='failed (y = 0)')
    ax.scatter(h[wrong], y[wrong], s=110, facecolors='none', edgecolors='#ff7b72', lw=1.5, zorder=5,
               label=f'misclassified by the line ({wrong.sum()})')
    ax.axvline(cut_line, color='#f0883e', ls='--', lw=1)
    ax.axvline(cut_log, color='#58a6ff', ls='--', lw=1)
    dx = xmax / 60
    ax.text(cut_line + dx, -0.52, f'line crosses 0.5\nat {cut_line:.2f} h', color='#f0883e', fontsize=7.5)
    if xmax < 20:
        ax.text(cut_log - dx, 1.3, f'logistic crosses 0.5\nat {cut_log:.2f} h', color='#58a6ff',
                fontsize=7.5, ha='right')
    else:   # no room on the left: write it to the right, with an arrow
        ax.annotate(f'logistic crosses 0.5 at {cut_log:.2f} h', xy=(cut_log, 1.3), xytext=(12, 1.3),
                    color='#58a6ff', fontsize=7.5, va='center',
                    arrowprops=dict(arrowstyle='->', color='#58a6ff', lw=0.8))
    ax.set_ylim(-0.6, 1.6)
    ax.set_xlim(0, xmax)
    ax.set_xlabel('hours of study')
    ax.set_title(title, fontsize=9.5, loc='left')
    ax.tick_params(labelsize=8)
axes[0].set_ylabel('prediction  (1 = pass)')
axes[0].text(8.2, 1.45, 'impossible (> 1)', fontsize=7.5, color='#ff7b72')
axes[0].text(0.4, -0.52, 'impossible (< 0)', fontsize=7.5, color='#ff7b72')
axes[1].legend(fontsize=7.5, loc='center right')
fig.suptitle('A line fitted to 0/1 labels leaves [0, 1] and is dragged by far, correct points; '
             'the logistic curve is not', fontsize=10)
plt.tight_layout()

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

plt.close()