import sklearn
from sklearn.datasets import make_classification
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import balanced_accuracy_score, classification_report, 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,
)

base_classifier = RandomForestClassifier(random_state=42)
base_classifier.fit(X_train, y_train)

default_predictions = base_classifier.predict(X_test)
default_score = balanced_accuracy_score(y_test, default_predictions)

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

tuned_predictions = tuned_classifier.predict(X_test)
tuned_score = balanced_accuracy_score(y_test, tuned_predictions)

print(f"scikit-learn {sklearn.__version__}")
print(f"default threshold balanced accuracy: {default_score:.3f}")
print(f"tuned threshold: {tuned_classifier.best_threshold_:.3f}")
print(f"cross-validated balanced accuracy: {tuned_classifier.best_score_:.3f}")
print(f"test balanced accuracy after tuning: {tuned_score:.3f}")
print()
print("Default threshold report")
print(classification_report(y_test, default_predictions, digits=3))
print("Tuned threshold report")
print(classification_report(y_test, tuned_predictions, digits=3))
print("Tuned threshold confusion matrix")
print(confusion_matrix(y_test, tuned_predictions))