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

# Why cross-entropy and not squared error, in three pictures
sigmoid = lambda z: 1 / (1 + np.exp(-z))
p = np.linspace(0.001, 0.999, 500)          # predicted probability for a point whose true class is 1

fig, axes = plt.subplots(1, 3, figsize=(10.5, 3.4))

ax = axes[0]
ax.plot(p, -np.log(p), color='#58a6ff', lw=2, label='cross-entropy  −ln p̂')
ax.plot(p, (1 - p) ** 2, color='#f0883e', lw=2, label='squared error  (1 − p̂)²')
ax.annotate(f'p̂ = 0.01: {-np.log(0.01):.1f}', xy=(0.01, -np.log(0.01)), xytext=(0.18, 4.4),
            fontsize=7.5, color='#58a6ff', arrowprops=dict(arrowstyle='->', color='gray', lw=0.8))
ax.annotate(f'p̂ = 0.01: {(1 - 0.01) ** 2:.2f}  (never above 1)', xy=(0.01, 0.98), xytext=(0.18, 1.7),
            fontsize=7.5, color='#f0883e', arrowprops=dict(arrowstyle='->', color='gray', lw=0.8))
ax.set_ylim(0, 5.5)
ax.set_xlabel('p̂ given to the true class (y = 1)')
ax.set_ylabel('loss of this one point')
ax.set_title('(a) the price of a mistake', fontsize=9.5, loc='left')
ax.legend(fontsize=7.5, loc='upper right')

ax = axes[1]
ax.plot(p, 1 - p, color='#58a6ff', lw=2, label='cross-entropy:  1 − p̂')
ax.plot(p, 2 * (1 - p) ** 2 * p, color='#f0883e', lw=2, label='squared error:  2(1 − p̂)²·p̂')
ax.annotate('confidently wrong:\nsquared error\nbarely pushes', xy=(0.02, 0.04), xytext=(0.42, 0.62),
            fontsize=7.5, color='#f0883e', arrowprops=dict(arrowstyle='->', color='gray', lw=0.8))
ax.set_ylim(0, 1.05)
ax.set_xlabel('p̂ given to the true class (y = 1)')
ax.set_ylabel('size of the gradient w.r.t. z')
ax.set_title('(b) how hard the gradient pushes', fontsize=9.5, loc='left')
ax.legend(fontsize=7.5, loc='upper right')

# (c) the whole loss as a function of one weight, on the 23 students of the lesson
hours = np.arange(1, 12.5, 0.5)
passed = np.array([0] * 11 + [1, 0] + [1] * 10)
ws = np.linspace(-3, 3, 400)
P = sigmoid(np.outer(ws, hours - 6.5))
mse = ((passed - P) ** 2).mean(axis=1)
ce = -(passed * np.log(P) + (1 - passed) * np.log(1 - P)).mean(axis=1)
ax = axes[2]
ax.plot(ws, ce, color='#58a6ff', lw=2, label='cross-entropy (left axis)')
ax.set_ylabel('cross-entropy', color='#58a6ff')
ax.set_ylim(0, 9)
ax.set_xlabel('weight w  (model p̂ = σ(w·(hours − 6.5)))')
ax.set_title('(c) the loss over one weight, 23 students', fontsize=9.5, loc='left')
axr = ax.twinx()
axr.plot(ws, mse, color='#f0883e', lw=2, label='mean squared error (right axis)')
axr.set_ylabel('mean squared error', color='#f0883e')
axr.set_ylim(0, 1)
axr.annotate('plateau: slope ≈ 0,\nconcave, not a bowl', xy=(-2.5, mse[np.searchsorted(ws, -2.5)]),
             xytext=(0.1, 0.72), fontsize=7.5, color='#f0883e',
             arrowprops=dict(arrowstyle='->', color='gray', lw=0.8))
ax.text(0.5, 3.2, 'cross-entropy:\nconvex, one minimum', color='#58a6ff', fontsize=7.5)
for a in list(axes) + [axr]:
    a.tick_params(labelsize=8)

plt.tight_layout()

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

plt.close()