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

# How CART picks a question: sort the customers by one feature, try a threshold at every midpoint
# between neighboring values, score each split by the weighted Gini of the two children, keep
# the lowest. Top: the 12 customers on the bill axis. Bottom: the weighted Gini of each candidate.
BLUE, ORANGE, GREEN, PURPLE = '#58a6ff', '#f0883e', '#3fb950', '#bc8cff'
bill = np.array([95, 120, 60, 110, 45, 130, 70, 115, 55, 125, 100, 65])
churn = np.array([1, 1, 0, 1, 0, 1, 0, 0, 0, 1, 0, 0])
names = np.array(list('ABCDEFGHIJKL'))


def gini(y):
    p = np.bincount(y, minlength=2) / len(y)
    return 1 - np.sum(p ** 2)


values = np.unique(bill)
mids = (values[:-1] + values[1:]) / 2
scores = np.array([(np.sum(bill <= t) * gini(churn[bill <= t]) + np.sum(bill > t) * gini(churn[bill > t]))
                   / len(bill) for t in mids])
best = mids[np.argmin(scores)]

plt.rcParams.update({'text.color': 'gray', 'axes.labelcolor': 'gray', 'axes.edgecolor': 'gray',
                     'xtick.color': 'gray', 'ytick.color': 'gray'})
fig, (top, bot) = plt.subplots(2, 1, figsize=(10, 5.6), sharex=True, gridspec_kw={'height_ratios': [1, 1.6]})
for t in mids:
    top.axvline(t, color=PURPLE if t == best else 'gray', lw=2.2 if t == best else 0.8,
                ls='-' if t == best else ':')
top.axvspan(40, best, color=BLUE, alpha=0.1, lw=0)
top.axvspan(best, 135, color=ORANGE, alpha=0.1, lw=0)
top.scatter(bill[churn == 0], np.zeros(np.sum(churn == 0)), s=90, color=BLUE, label='stayed', zorder=3)
top.scatter(bill[churn == 1], np.zeros(np.sum(churn == 1)), s=110, color=ORANGE, marker='^',
            label='churned', zorder=3)
for b, n in zip(bill, names):
    top.text(b, -0.32, n, ha='center', va='center', fontsize=9)
top.text(61, -0.85, 'left child: 5 stayed · 0 churned\nGini 0', ha='center', fontsize=9, color=BLUE,
         fontweight='bold')
top.text(110, -0.85, 'right child: 2 stayed · 5 churned\nGini 0.408', ha='center', fontsize=9, color=ORANGE,
         fontweight='bold')
top.set_ylim(-1.2, 0.75); top.set_yticks([])
top.legend(loc='upper center', fontsize=8.5, frameon=False, ncol=2)
top.set_title(f'11 candidate thresholds (dotted) between 12 bills; the winner, bill ≤ {best:g}, in purple',
              fontsize=10)
for s in ['top', 'right', 'left']:
    top.spines[s].set_visible(False)

colors = [PURPLE if t == best else 'gray' for t in mids]
bot.bar(mids, scores, width=3.2, color=colors, alpha=0.75)
bot.axhline(gini(churn), color='gray', ls='--', lw=1)
bot.text(134, gini(churn) + 0.01, f'root Gini {gini(churn):.3f}', ha='right', va='bottom', fontsize=8.5)
for t, s in zip(mids, scores):
    bot.text(t, s + 0.008, f'{s:.3f}', ha='center', fontsize=7.5, color=PURPLE if t == best else 'gray',
             fontweight='bold' if t == best else 'normal')
bot.set_ylim(0, 0.55)
bot.set_ylabel('weighted Gini of\nthe two children')
bot.set_xlabel('monthly bill (R$) — each bar is one candidate question "bill ≤ t?"')
bot.set_xlim(40, 135)
for s in ['top', 'right']:
    bot.spines[s].set_visible(False)
plt.tight_layout()

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

plt.close()