# Source: content/notes/ml/trees-and-ensembles.md
# Independent CPU example; use the curriculum environment.
# See /notes/ml/#example-environment or /notes/deep-learning/#example-environment.

import numpy as np
from sklearn.datasets import make_classification
from sklearn.metrics import balanced_accuracy_score, log_loss
from sklearn.model_selection import GridSearchCV, StratifiedKFold, train_test_split
from sklearn.tree import DecisionTreeClassifier, export_text

X, y = make_classification(n_samples=600, n_features=8, n_informative=5,
                           n_redundant=1, class_sep=1.2, random_state=21)
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.25, stratify=y, random_state=21)
search = GridSearchCV(DecisionTreeClassifier(random_state=21),
    {"max_depth": [3, 6, None], "min_samples_leaf": [5, 15],
     "ccp_alpha": [0.0, 0.005]}, scoring="neg_log_loss",
    cv=StratifiedKFold(3, shuffle=True, random_state=21), n_jobs=1)
search.fit(X_train, y_train)
tree = search.best_estimator_
pred = tree.predict(X_test)
prob = tree.predict_proba(X_test)
assert np.allclose(prob.sum(axis=1), 1)
assert tree.get_n_leaves() <= len(X_train)
assert balanced_accuracy_score(y_test, pred) > 0.6
print(search.best_params_)
print("Balanced accuracy:", balanced_accuracy_score(y_test, pred))
print("Log loss:", log_loss(y_test, prob))
print(export_text(tree, max_depth=2))