import os
import sys

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

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from _treedraw import draw_tree  # noqa: E402

# Growing the lesson's churn tree one level at a time. Each level re-runs the same search
# (the best question by weighted Gini) inside every child that is still mixed: greedy and
# recursive. Under each panel, the number of leaves and the training accuracy.
BLUE, ORANGE = '#58a6ff', '#f0883e'
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]

plt.rcParams.update({'text.color': 'gray'})
fig, axes = plt.subplots(2, 2, figsize=(10.5, 8.4))
titles = ['depth 0: no question yet', 'depth 1: the best first question',
          'depth 2: split the mixed child again', 'depth 3: until every leaf is pure']
for d, (ax, title) in enumerate(zip(axes.flat, titles)):
    if d == 0:
        ax.text(0, 0, '→ stay\n7 · 5\nGini 0.486', ha='center', va='center', fontsize=9, color='gray',
                linespacing=1.35, bbox=dict(boxstyle='round,pad=0.4', fc=BLUE, ec='gray', alpha=0.2))
        leaves, acc = 1, np.mean(churn == 0)
        xs = [0]
    else:
        tree = DecisionTreeClassifier(max_depth=d, random_state=0).fit(X, churn)
        xy = draw_tree(ax, tree, ['tenure', 'bill'], ['stay', 'churn'], [BLUE, ORANGE],
                       xscale=1.5, yscale=1.25, fontsize=8.5)
        leaves, acc = tree.get_n_leaves(), tree.score(X, churn)
        xs = [x for x, _ in xy.values()]
    ax.set_xlim(min(xs) - 1.1, max(xs) + 1.1)
    ax.set_ylim(-3 * 1.25 - 1.1, 0.7)
    ax.axis('off')
    ax.set_title(title, fontsize=10.5, fontweight='bold', color='gray')
    ax.text(np.mean([min(xs), max(xs)]), -3 * 1.25 - 0.95,
            f'{leaves} {"leaves" if leaves > 1 else "leaf"} · training accuracy {acc * 12:.0f}/12 = {acc:.2f}',
            ha='center', fontsize=9.5, color='gray')
fig.text(0.5, 0.005, 'numbers in each box: stayed · churned.  The depth-3 split exists to isolate a single '
         'customer — 12/12 on training, nothing learned.', ha='center', fontsize=9, style='italic', color='gray')
plt.tight_layout(rect=(0, 0.03, 1, 1), h_pad=2)

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

plt.close()