import numpy as np
from sklearn.base import clone
from sklearn.datasets import make_moons
from sklearn.dummy import DummyClassifier
from sklearn.ensemble import RandomForestClassifier
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import (balanced_accuracy_score, confusion_matrix,
                             f1_score, roc_auc_score)
from sklearn.model_selection import (StratifiedKFold, cross_validate,
                                     train_test_split)
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler
from sklearn.tree import DecisionTreeClassifier

X, y = make_moons(n_samples=1000, noise=0.25, random_state=7)
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42, stratify=y
)
cv = StratifiedKFold(n_splits=5, shuffle=True, random_state=42)
folds = list(cv.split(X_train, y_train))
models = {
    "dummy": DummyClassifier(strategy="most_frequent"),
    "logistic": make_pipeline(StandardScaler(),
                              LogisticRegression(max_iter=1000)),
    "tree": DecisionTreeClassifier(max_depth=4, random_state=42),
    "forest": RandomForestClassifier(n_estimators=100,
                                    min_samples_leaf=3,
                                    random_state=42, n_jobs=1),
}
means = {}
for name, estimator in models.items():
    scores = cross_validate(estimator, X_train, y_train, cv=folds,
                            scoring="balanced_accuracy",
                            return_train_score=True)
    means[name] = scores["test_score"].mean()
    # Here test_score means each CV validation fold, not X_test.
    print(name, "train", round(scores["train_score"].mean(), 3),
          "validation", round(means[name], 3),
          "fold std", round(scores["test_score"].std(), 3))

winner_name = max(means, key=means.get)
winner = clone(models[winner_name]).fit(X_train, y_train)
prediction = winner.predict(X_test)
positive_column = np.flatnonzero(winner.classes_ == 1)[0]
probability = winner.predict_proba(X_test)[:, positive_column]
assert prediction.shape == y_test.shape
assert np.isfinite(probability).all()
print("Selected:", winner_name)
print("Test balanced accuracy:", balanced_accuracy_score(y_test, prediction))
print("Test F1:", f1_score(y_test, prediction))
print("Test ROC AUC:", roc_auc_score(y_test, probability))
print("Confusion matrix; rows=true, columns=predicted, classes=[0, 1]:")
print(confusion_matrix(y_test, prediction, labels=[0, 1]))