# %%
import numpy as np
from sklearn.datasets import make_classification

X, y = make_classification(
    n_samples=10_000, weights=[0.1, 0.9], random_state=42, class_sep=1
)
sample_weight = np.zeros_like(y, dtype=np.float64)
sample_weight[y == 0] = 0.1
sample_weight[y == 1] = 0.9

# %%
from sklearn.model_selection import train_test_split

X_train, X_test, y_train, y_test, sw_train, sw_test = train_test_split(
    X, y, sample_weight, random_state=42
)

# %%
from sklearn.linear_model import LogisticRegression
from sklearn.calibration import CalibratedClassifierCV

calibration_method = "isotonic"
models = {
    "LR no weights": LogisticRegression(),
    "LR class weights": LogisticRegression(class_weight="balanced"),
    "Calibrated LR no weights": CalibratedClassifierCV(
        LogisticRegression(),
        method=calibration_method,
    ),
    "Calibrated LR class weights": CalibratedClassifierCV(
        LogisticRegression(class_weight="balanced"),
        method=calibration_method,
    ),
    "Calibrated LR sample weights": CalibratedClassifierCV(
        LogisticRegression(),
        method=calibration_method,
    ),
    "Calibrated LR class and sample weights": CalibratedClassifierCV(
        LogisticRegression(class_weight="balanced"),
        method=calibration_method,
    ),
}

# %%
import matplotlib.pyplot as plt
from sklearn.calibration import CalibrationDisplay
from sklearn.metrics import balanced_accuracy_score

fig, ax = plt.subplots()

calibration_display_params = {
    "n_bins": 20,
    "strategy": "quantile",
}
for name, model in models.items():
    if "sample weights" in name:
        model.fit(X_train, y_train, sample_weight=sw_train)
    else:
        model.fit(X_train, y_train)

    score = balanced_accuracy_score(y_test, model.predict(X_test))
    CalibrationDisplay.from_estimator(
        model,
        X_test,
        y_test,
        name=name + f" - {score:.3f}",
        ax=ax,
        **calibration_display_params,
    )
ax.legend(loc="center left", bbox_to_anchor=(1, 0.5), title="Model - Balanced Accuracy")
_ = fig.suptitle(f"Using {calibration_method} calibration")

# %%
from sklearn.model_selection import cross_validate
from sklearn.model_selection import KFold

cv_results = {}
cv = KFold(n_splits=5)
for name, model in models.items():
    if "sample weights" in name:
        fit_params = {"sample_weight": sample_weight}
    else:
        fit_params = {}
    cv_results[name] = cross_validate(
        model,
        X,
        y,
        cv=cv,
        fit_params=fit_params,
        scoring="balanced_accuracy",
        return_estimator=True,
        return_indices=True,
    )

# %%
fig, ax = plt.subplots()
for model_idx, (name, results) in enumerate(cv_results.items()):
    CalibrationDisplay.from_cv_results(
        results, X, y, ax=ax, name=name, **calibration_display_params
    )
ax.legend(loc="center left", bbox_to_anchor=(1, 0.5))

# %%
fig, ax = plt.subplots()
for model_idx, (name, results) in enumerate(cv_results.items()):
    CalibrationDisplay.from_cv_results(
        results,
        X,
        y,
        ax=ax,
        name=name,
        plot_uncertainty_style="fill_between",
        **calibration_display_params,
    )
ax.legend(loc="center left", bbox_to_anchor=(1, 0.5))

# %%
fig, ax = plt.subplots()
for model_idx, (name, results) in enumerate(cv_results.items()):
    CalibrationDisplay.from_cv_results(
        results,
        X,
        y,
        ax=ax,
        name=name,
        plot_uncertainty_style="lines",
        **calibration_display_params,
    )
ax.legend(loc="center left", bbox_to_anchor=(1, 0.5))

# %%