import sklearn
from sklearn.datasets import make_classification
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import balanced_accuracy_score, confusion_matrix
from sklearn.model_selection import TunedThresholdClassifierCV, train_test_split

X, y = make_classification(
    n_samples=1_200,
    n_features=12,
    n_informative=5,
    weights=[0.86, 0.14],
    class_sep=0.7,
    random_state=42,
)

X_train, X_test, y_train, y_test = train_test_split(
    X,
    y,
    stratify=y,
    test_size=0.25,
    random_state=42,
)

default_classifier = RandomForestClassifier(random_state=42)
default_classifier.fit(X_train, y_train)
default_predictions = default_classifier.predict(X_test)

tuned_classifier = TunedThresholdClassifierCV(
    estimator=RandomForestClassifier(random_state=42),
    scoring="balanced_accuracy",
    thresholds=50,
    cv=5,
)
tuned_classifier.fit(X_train, y_train)
tuned_predictions = tuned_classifier.predict(X_test)

print(f"scikit-learn {sklearn.__version__}")
print(f"selected threshold: {tuned_classifier.best_threshold_:.3f}")
print(
    "cross-validated balanced accuracy: "
    f"{tuned_classifier.best_score_:.3f}"
)
print(
    "default test balanced accuracy: "
    f"{balanced_accuracy_score(y_test, default_predictions):.3f}"
)
print(
    "tuned test balanced accuracy: "
    f"{balanced_accuracy_score(y_test, tuned_predictions):.3f}"
)
print("default confusion matrix:")
print(confusion_matrix(y_test, default_predictions))
print("tuned confusion matrix:")
print(confusion_matrix(y_test, tuned_predictions))