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

# The two branches of the log loss, with the four hand-worked predictions of the lesson on them
p = np.linspace(0.005, 0.995, 500)          # predicted probability of class 1, p̂

fig, ax = plt.subplots(figsize=(7.5, 3.6))
ax.plot(p, -np.log(p), color='#58a6ff', lw=2, label='cost when y = 1:  −ln p̂')
ax.plot(p, -np.log(1 - p), color='#f0883e', lw=2, label='cost when y = 0:  −ln(1 − p̂)')

# the table of the lesson: (y, p̂)
points = [(1, 0.9), (1, 0.6), (0, 0.3), (0, 0.8)]
label_at = [(0.80, 0.55), (0.55, 0.12), (0.06, 0.62), (0.52, 2.1)]    # where each label is written
for i, ((y, ph), xy_text) in enumerate(zip(points, label_at), start=1):
    cost = -np.log(ph) if y == 1 else -np.log(1 - ph)
    color = '#58a6ff' if y == 1 else '#f0883e'
    ax.plot(ph, cost, 'o', color=color, ms=7, mec='black', mew=0.8, zorder=3)
    ax.annotate(f'{i}: y = {y}, p̂ = {ph} → {cost:.3f}', xy=(ph, cost), xytext=xy_text,
                fontsize=7.5, color=color, arrowprops=dict(arrowstyle='->', color='gray', lw=0.7))

ax.annotate('p̂ = 0.5 costs ln 2 = 0.693\nfor either class', xy=(0.5, np.log(2)), xytext=(0.36, 1.45),
            fontsize=7.5, color='gray', ha='center', arrowprops=dict(arrowstyle='->', color='gray', lw=0.7))
ax.set_ylim(0, 5.3)
ax.set_xlim(0, 1)
ax.set_xlabel('p̂  (predicted probability of class 1)')
ax.set_ylabel('cost of one example')
ax.legend(fontsize=8, loc='upper center')
ax.tick_params(labelsize=8)
ax.grid(alpha=0.3)

plt.tight_layout()

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

plt.close()