# 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.ensemble import AdaBoostClassifier
from sklearn.metrics import balanced_accuracy_score, log_loss
from sklearn.model_selection import train_test_split
from sklearn.tree import DecisionTreeClassifier

X, y = make_classification(n_samples=500, n_features=6, n_informative=4,
                           n_redundant=0, class_sep=1.4, random_state=5)
X_train, X_test, y_train, y_test = train_test_split(
    X, y, stratify=y, random_state=5)
model = AdaBoostClassifier(
    estimator=DecisionTreeClassifier(max_depth=1, random_state=5),
    n_estimators=70, learning_rate=0.5, random_state=5)
# Older releases expose this selector; newer releases use SAMME directly.
if "algorithm" in model.get_params():
    model.set_params(algorithm="SAMME")
model.fit(X_train, y_train)
pred = model.predict(X_test)
prob = model.predict_proba(X_test)
assert len(model.estimators_) > 0
assert np.allclose(prob.sum(axis=1), 1)
print("Balanced accuracy:", balanced_accuracy_score(y_test, pred))
print("Log loss:", log_loss(y_test, prob))
print("First stage errors:", model.estimator_errors_[:5])