import numpy as np
import sklearn
from sklearn.calibration import CalibratedClassifierCV, calibration_curve
from sklearn.datasets import make_classification
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import brier_score_loss, log_loss, roc_auc_score
from sklearn.model_selection import train_test_split


X, y = make_classification(
    n_samples=4_000,
    n_features=20,
    n_informative=7,
    n_redundant=5,
    weights=[0.78, 0.22],
    class_sep=0.65,
    flip_y=0.08,
    random_state=42,
)

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

base_classifier = RandomForestClassifier(
    n_estimators=200,
    min_samples_leaf=8,
    random_state=42,
)
base_classifier.fit(X_train, y_train)
base_probability = base_classifier.predict_proba(X_test)[:, 1]

calibrated_classifier = CalibratedClassifierCV(
    estimator=RandomForestClassifier(
        n_estimators=200,
        min_samples_leaf=8,
        random_state=42,
    ),
    method="sigmoid",
    cv=5,
)
calibrated_classifier.fit(X_train, y_train)
calibrated_probability = calibrated_classifier.predict_proba(X_test)[:, 1]


def print_scores(name, probability):
    print(f"{name} Brier loss: {brier_score_loss(y_test, probability):.4f}")
    print(f"{name} log loss: {log_loss(y_test, probability):.4f}")
    print(f"{name} ROC AUC: {roc_auc_score(y_test, probability):.4f}")


print(f"scikit-learn {sklearn.__version__}")
print(f"training rows: {X_train.shape[0]}")
print(f"test rows: {X_test.shape[0]}")
print(f"calibration folds: {len(calibrated_classifier.calibrated_classifiers_)}")
print()
print_scores("uncalibrated", base_probability)
print_scores("calibrated", calibrated_probability)
print()
print("calibrated probability bins:")
fraction_positive, mean_predicted = calibration_curve(
    y_test,
    calibrated_probability,
    n_bins=5,
    strategy="quantile",
)

for bin_number, (predicted, observed) in enumerate(
    zip(mean_predicted, fraction_positive),
    start=1,
):
    print(
        f"bin {bin_number}: "
        f"mean predicted={predicted:.3f}, "
        f"observed positive rate={observed:.3f}"
    )

print()
print("first five calibrated probabilities:")
print(np.array2string(calibrated_probability[:5], precision=3, floatmode="fixed"))