# 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 catboost import CatBoostClassifier
from sklearn.metrics import balanced_accuracy_score, log_loss
from sklearn.model_selection import train_test_split

rng = np.random.default_rng(31)
numeric = rng.normal(size=(600, 2))
region = rng.choice(["north", "south", "east"], size=600)
score = numeric[:, 0] - 0.7 * numeric[:, 1] + 0.9 * (region == "north")
y = rng.binomial(1, 1 / (1 + np.exp(-score)))
X = np.empty((600, 3), dtype=object)
X[:, :2] = numeric
X[:, 2] = region
dev, test = train_test_split(np.arange(600), test_size=0.2,
                              stratify=y, random_state=31)
train, valid = train_test_split(dev, test_size=0.25,
                                stratify=y[dev], random_state=32)
model = CatBoostClassifier(iterations=140, depth=4, learning_rate=0.06,
    loss_function="Logloss", cat_features=[2], random_seed=31,
    thread_count=1, task_type="CPU", allow_writing_files=False, verbose=False)
model.fit(X[train], y[train], eval_set=(X[valid], y[valid]),
          early_stopping_rounds=12, use_best_model=True)
pred = np.asarray(model.predict(X[test])).reshape(-1)
prob = model.predict_proba(X[test])[:, 1]
assert pred.shape == y[test].shape
assert np.isfinite(prob).all() and 0 < model.tree_count_ <= 140
print("Trees retained:", model.tree_count_)
print("Test log loss:", log_loss(y[test], prob))
print("Balanced accuracy:", balanced_accuracy_score(y[test], pred))