import matplotlib.pyplot as plt
import numpy as np
from io import StringIO
from matplotlib.patches import Rectangle
from sklearn.tree import DecisionTreeClassifier

# The lesson's 12 customers: the depth-2 tree CART grows on them, drawn twice.
# Left: the partition of the plane into axis-aligned rectangles.
# Right: the same tree as a flowchart, with the class counts of every node.
# Leaves are numbered by their node id, the number tree.apply() returns.
tenure = np.array([2, 3, 5, 6, 8, 11, 14, 19, 24, 30, 38, 50])
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])
X = np.c_[tenure, bill]
tree = DecisionTreeClassifier(max_depth=2, random_state=0).fit(X, churn)
t = tree.tree_
assert t.feature[0] == 1 and t.feature[2] == 0      # root on bill, then tenure
b0, t0 = t.threshold[0], t.threshold[2]

BLUE, ORANGE = '#58a6ff', '#f0883e'
plt.rcParams.update({'text.color': 'gray', 'axes.labelcolor': 'gray', 'axes.edgecolor': 'gray',
                     'xtick.color': 'gray', 'ytick.color': 'gray'})
fig, (ax, ax2) = plt.subplots(1, 2, figsize=(10, 4.2), gridspec_kw={'width_ratios': [1, 1.15]})

x0, x1, y0, y1 = 35, 140, 0, 55
ax.add_patch(Rectangle((x0, y0), b0 - x0, y1 - y0, color=BLUE, alpha=0.13, lw=0))
ax.add_patch(Rectangle((b0, y0), x1 - b0, t0 - y0, color=ORANGE, alpha=0.18, lw=0))
ax.add_patch(Rectangle((b0, t0), x1 - b0, y1 - t0, color=BLUE, alpha=0.13, lw=0))
ax.axvline(b0, color='gray', lw=1.6)
ax.plot([b0, x1], [t0, t0], color='gray', lw=1.6, ls='--')
ax.scatter(bill[churn == 0], tenure[churn == 0], s=46, color=BLUE, label='stayed', zorder=3)
ax.scatter(bill[churn == 1], tenure[churn == 1], s=46, color=ORANGE, marker='^', label='churned', zorder=3)
ax.text(b0 + 1, y1 - 1.5, f'bill = {b0:g}', fontsize=8, va='top')
ax.text(x1 - 1, t0 + 0.8, f'tenure = {t0:g}', fontsize=8, ha='right')
ax.text(69, 44, 'leaf 1\n5 stay · 0 churn', fontsize=8, ha='center', va='top')
ax.text(104, 12, 'leaf 3\n0 stay · 4 churn', fontsize=8, ha='center', va='center')
ax.text(90, 42, 'leaf 4\n2 stay · 1 churn', fontsize=8, ha='left', va='center')
ax.set_xlim(x0, x1); ax.set_ylim(y0, y1)
ax.set_xlabel('monthly bill (R$)')
ax.set_ylabel('tenure (months as a client)')
ax.set_title('Two questions cut the plane into 3 rectangles', fontsize=10)
ax.legend(fontsize=8, loc='upper right', frameon=False)


def box(x, y, text, color):
    ax2.text(x, y, text, ha='center', va='center', fontsize=8.5, color='gray',
             bbox=dict(boxstyle='round,pad=0.45', fc=color, ec='gray', alpha=0.25 if color != 'none' else 1))


def edge(xa, ya, xb, yb, label):
    ax2.annotate('', xy=(xb, yb + 0.07), xytext=(xa, ya - 0.07),
                 arrowprops=dict(arrowstyle='->', color='gray', lw=1.2))
    ax2.text((xa + xb) / 2 + (0.03 if xb > xa else -0.03), (ya + yb) / 2, label, fontsize=8,
             ha='left' if xb > xa else 'right', va='center')


def counts(node):
    v = t.value[node][0] * t.n_node_samples[node]
    return f"{v[0]:.0f} stay · {v[1]:.0f} churn\nGini {t.impurity[node]:.3f}"


box(0.5, 0.88, f'bill ≤ {b0:g}?\n' + counts(0), 'none')
box(0.2, 0.52, 'leaf 1 → stay\n' + counts(1), BLUE)
box(0.72, 0.52, f'tenure ≤ {t0:g}?\n' + counts(2), 'none')
box(0.5, 0.13, 'leaf 3 → churn\n' + counts(3), ORANGE)
p = t.value[4][0][1]
box(0.9, 0.13, f'leaf 4 → stay\n' + counts(4) + f'\nP(churn) = {p:.2f}', BLUE)
edge(0.46, 0.88, 0.22, 0.52, 'yes')
edge(0.54, 0.88, 0.7, 0.52, 'no')
edge(0.68, 0.52, 0.52, 0.13, 'yes')
edge(0.76, 0.52, 0.88, 0.13, 'no')
ax2.set_xlim(0, 1.08); ax2.set_ylim(0, 1)
ax2.axis('off')
ax2.set_title('The same tree as a flowchart', fontsize=10)

plt.tight_layout()

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

plt.close()