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

# Left: the three impurity measures of a two-class node, as a function of the
# share p of the positive class. Right: why misclassification rate is a poor
# split criterion. Parent node: 10 rows, 2 positive (p = 0.2). The split sends
# 5 negatives to a pure child (p = 0) and leaves 3 negatives + 2 positives in
# the other (p = 0.4). The weighted impurity of the children is the midpoint
# of the chord between the two children's points on each curve.
p = np.linspace(0, 1, 401)
gini = 1 - p**2 - (1 - p)**2
with np.errstate(divide='ignore', invalid='ignore'):
    ent = -(np.where(p > 0, p * np.log2(p), 0) + np.where(p < 1, (1 - p) * np.log2(1 - p), 0))
err = np.minimum(p, 1 - p)

BLUE, ORANGE, GREEN, RED = '#58a6ff', '#f0883e', '#3fb950', '#ff7b72'
plt.rcParams.update({'text.color': 'gray', 'axes.labelcolor': 'gray', 'axes.edgecolor': 'gray',
                     'xtick.color': 'gray', 'ytick.color': 'gray'})
fig, axes = plt.subplots(1, 2, figsize=(9.6, 3.8))

ax = axes[0]
ax.plot(p, ent, color=ORANGE, lw=2, label='entropy  $-\\sum p_c \\log_2 p_c$')
ax.plot(p, gini, color=BLUE, lw=2, label='Gini  $1-\\sum p_c^2$')
ax.plot(p, err, color=RED, lw=2, label='misclassification  $1-\\max_c p_c$')
ax.set_xlabel('share of the positive class in the node, p')
ax.set_ylabel('impurity')
ax.set_title('Three ways to score a node (two classes)', fontsize=10)
ax.legend(fontsize=7.5, loc='upper center', frameon=False)
ax.set_ylim(0, 1.42)
ax.set_yticks([0, 0.2, 0.4, 0.6, 0.8, 1.0])
ax.text(0.5, 0.03, 'all three: 0 when pure,\nmaximum at p = 0.5', ha='center', fontsize=7.5, color='gray')

ax = axes[1]
pp, pl, pr, wl, wr = 0.2, 0.0, 0.4, 0.5, 0.5          # parent, children, child weights
g = lambda q: 1 - q**2 - (1 - q)**2
e = lambda q: min(q, 1 - q)
ax.plot(p[p <= 0.5], gini[p <= 0.5], color=BLUE, lw=2, label='Gini')
ax.plot(p[p <= 0.5], err[p <= 0.5], color=RED, lw=2, label='misclassification')
for f, c in [(g, BLUE), (e, RED)]:
    ax.plot([pl, pr], [f(pl), f(pr)], '--', color=c, lw=1)            # chord between the children
    ax.scatter([pl, pr], [f(pl), f(pr)], color=c, s=28, zorder=5)
    ax.scatter([pp], [f(pp)], color=c, s=60, marker='s', zorder=5)       # parent on the curve
    ax.scatter([pp], [wl * f(pl) + wr * f(pr)], color=c, s=60, marker='o',
               facecolors='none', lw=1.6, zorder=6)                       # children, weighted
ax.annotate('', xy=(pp, wl * g(pl) + wr * g(pr) + 0.004), xytext=(pp, g(pp) - 0.004),
            arrowprops=dict(arrowstyle='->', color=GREEN, lw=1.6))
ax.text(pp + 0.012, 0.285, f'Gini drops\n{g(pp):.2f} → {wl * g(pl) + wr * g(pr):.2f}', fontsize=7.5, color=GREEN)
ax.text(pp + 0.012, 0.155, f'error stays\n{e(pp):.2f} → {wl * e(pl) + wr * e(pr):.2f}', fontsize=7.5, color=RED)
ax.text(0.005, 0.43, 'parent (■): 8 neg + 2 pos\nleft child: 5 neg + 0 pos\nright child: 3 neg + 2 pos',
        fontsize=7.5, color='gray', va='top')
ax.set_xlabel('share of the positive class, p')
ax.set_ylabel('impurity')
ax.set_xlim(-0.02, 0.52)
ax.set_ylim(-0.02, 0.52)
ax.set_title('A useful split the error rate cannot see', fontsize=10)
ax.legend(fontsize=7.5, loc='lower right', frameon=False)

for a in axes:
    a.spines[['top', 'right']].set_visible(False)
plt.tight_layout()

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

plt.close()