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

# Anatomy of a decision tree: the lesson's 12 customers, grown until every leaf is pure.
# Labels name the parts (root, internal node, branch, leaf, depth) and the highlighted path is
# how a new customer (25 months, R$ 110) travels from the root to the leaf that answers.
BLUE, ORANGE, PURPLE = '#58a6ff', '#f0883e', '#bc8cff'
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(random_state=0).fit(X, churn)
t = tree.tree_
new = np.array([[25, 110]])
path = list(tree.decision_path(new).indices)

plt.rcParams.update({'text.color': 'gray'})
fig, ax = plt.subplots(figsize=(10, 6.2))
xy = draw_tree(ax, tree, ['tenure', 'bill'], ['stay', 'churn'], [BLUE, ORANGE], xscale=2.3, yscale=1.55,
               x0=1.6, y0=0, path=path, path_color=PURPLE, fontsize=9,
               leaf_text=lambda n, c: f"leaf → {'stay' if c[0] > c[1] else 'churn'}")
leaves = [n for n in range(t.node_count) if t.children_left[n] == -1]
single = next(n for n in leaves if t.n_node_samples[n] == 1)     # customer J, alone in a leaf
xmin = min(x for x, _ in xy.values()) - 1.7

for d in range(tree.get_depth() + 1):
    ax.text(xmin, -d * 1.55, f'depth {d}', fontsize=9.5, ha='left', va='center', fontweight='bold')
    ax.plot([xmin + 0.85, xmin + 1.1], [-d * 1.55] * 2, color='gray', lw=0.8, ls=':')


def callout(node, text, dx, dy):
    x, y = xy[node]
    ax.annotate(text, xy=(x + np.sign(dx) * 0.55, y + np.sign(dy) * 0.1), xytext=(x + dx, y + dy),
                fontsize=9, color=PURPLE if 'path' in text else 'gray', fontweight='bold', ha='center',
                arrowprops=dict(arrowstyle='->', color='gray', lw=0.9))


callout(0, 'ROOT: the first question,\nasked of every customer', 2.6, 0.25)
callout(2, 'INTERNAL NODE:\nanother question', 2.9, 0.15)
callout(leaves[0], 'LEAF: no question,\nonly an answer', -0.3, -1.35)
callout(single, 'a leaf holding ONE customer (J):\nmemorized, not learned', 2.0, 0.85)
ax.text(xy[0][0] - 2.9, -0.55, 'BRANCH: "yes" goes left,\n"no" goes right', fontsize=9, fontweight='bold',
        ha='center')
ax.text(xmin, 0.65, 'numbers in each box: customers who stayed · churned', fontsize=8.5, style='italic')
ax.text(xmin, -tree.get_depth() * 1.55 - 1.1,
        'highlighted path: a new customer, 25 months, R$ 110 → bill > 82.5 → tenure > 15 → bill ≤ 120 → stay',
        fontsize=9.5, color=PURPLE, fontweight='bold')
ax.set_xlim(xmin - 0.1, max(x for x, _ in xy.values()) + 3.4)
ax.set_ylim(-tree.get_depth() * 1.55 - 1.4, 0.9)
ax.axis('off')
plt.tight_layout()

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

plt.close()