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

# Weakest-link pruning on the lesson's fully grown churn tree. Every internal node t shows
# alpha_eff(t) = (R(t) - R(T_t)) / (|T_t| - 1): the price per leaf saved by collapsing its subtree.
# Each step collapses the node with the smallest alpha_eff (outlined in purple) and recomputes.
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_
n = t.weighted_n_node_samples
R = t.impurity * n / n[0]                    # each node's impurity, weighted by its share of rows


def subtree(node):
    if t.children_left[node] == -1:
        return [node]
    return [node] + subtree(t.children_left[node]) + subtree(t.children_right[node])


def leaves_under(node, collapsed):
    if t.children_left[node] == -1 or node in collapsed:
        return [node]
    return leaves_under(t.children_left[node], collapsed) + leaves_under(t.children_right[node], collapsed)


def internal(collapsed, hidden):
    return [k for k in range(t.node_count) if t.children_left[k] != -1 and k not in collapsed and k not in hidden]


plt.rcParams.update({'text.color': 'gray'})
fig, axes = plt.subplots(1, 4, figsize=(15, 4.9))
collapsed, hidden, alpha_now = set(), set(), 0.0
for step, ax in enumerate(axes):
    alive = internal(collapsed, hidden)
    eff = {k: (R[k] - sum(R[l] for l in leaves_under(k, collapsed))) / (len(leaves_under(k, collapsed)) - 1)
           for k in alive}
    weakest = min(eff, key=eff.get) if eff else None
    xy = draw_tree(ax, tree, ['tenure', 'bill'], ['stay', 'churn'], [BLUE, ORANGE], xscale=1.35, yscale=1.3,
                   fontsize=8, show_gini=False, collapsed=collapsed, hidden=hidden,
                   path=[weakest] if weakest is not None else [], path_color=PURPLE,
                   node_labels={k: f'α_eff {v:.3f}' for k, v in eff.items()})
    n_leaves = len(leaves_under(0, collapsed))
    ax.set_title(f'α = {alpha_now:.3f}: {n_leaves} {"leaves" if n_leaves > 1 else "leaf"}', fontsize=10.5,
                 fontweight='bold', color='gray')
    xs = [x for x, _ in xy.values()]
    ax.set_xlim(min(xs) - 0.8, max(xs) + 0.8); ax.set_ylim(-3 * 1.3 - 0.6, 0.6)
    ax.axis('off')
    if weakest is not None:
        ax.text(np.mean([min(xs), max(xs)]), -3 * 1.3 - 0.45, f'next: collapse the purple node at α = {eff[weakest]:.3f}',
                ha='center', fontsize=8.5, color=PURPLE)
        alpha_now = eff[weakest]
        collapsed.add(weakest)
        hidden.update(set(subtree(weakest)) - {weakest})
fig.text(0.5, 0.01, 'numbers in each box: stayed · churned.  The sequence 0.111 → 0.127 → 0.248 is the '
         'one cost_complexity_pruning_path returns.', ha='center', fontsize=9, style='italic', color='gray')
plt.tight_layout(rect=(0, 0.04, 1, 1))

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

plt.close()