# Source: content/notes/ml/trees-and-ensembles.md
# Independent CPU example; use the curriculum environment plus XGBoost/CatBoost.
# 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 log_loss, roc_auc_score
from sklearn.model_selection import train_test_split
from xgboost import XGBClassifier

X, y = make_classification(n_samples=750, n_features=10, n_informative=6,
                           n_redundant=2, random_state=23)
X_dev, X_test, y_dev, y_test = train_test_split(
    X, y, test_size=0.2, stratify=y, random_state=23)
X_train, X_valid, y_train, y_valid = train_test_split(
    X_dev, y_dev, test_size=0.25, stratify=y_dev, random_state=24)
model = XGBClassifier(n_estimators=150, max_depth=3, learning_rate=0.08,
    tree_method="hist", objective="binary:logistic", eval_metric="logloss",
    early_stopping_rounds=12, subsample=0.9, colsample_bytree=0.9,
    random_state=23, n_jobs=1)
model.fit(X_train, y_train, eval_set=[(X_valid, y_valid)], verbose=False)
pred = model.predict(X_test)
prob = model.predict_proba(X_test)[:, 1]
assert pred.shape == y_test.shape
assert np.isfinite(prob).all() and ((prob >= 0) & (prob <= 1)).all()
assert 0 <= model.best_iteration < 150
print("Best iteration:", model.best_iteration)
print("Test log loss:", log_loss(y_test, prob))
print("Test ROC AUC:", roc_auc_score(y_test, prob))