import numpy as np
import matplotlib.pyplot as plt
from sklearn.linear_model import LogisticRegression
from sklearn.inspection import DecisionBoundaryDisplay


data = np.array(
    [
        [-1, -1],
        [-2, -1],
        [1, 1],
        [2, 1],
        [2, 2],
        [3, 2],
        [3, 3],
    ]
)
target = np.arange(data.shape[0])
clf = LogisticRegression().fit(data, target)

cmap = "gist_rainbow"
_, axes = plt.subplots(nrows=3, ncols=3, figsize=(12, 12), constrained_layout=True)
for plot_method_idx, plot_method in enumerate(["contourf", "contour", "pcolormesh"]):
    for response_method_idx, response_method in enumerate(
        ["predict_proba", "decision_function", "predict"]
    ):
        ax = axes[plot_method_idx, response_method_idx]
        display = DecisionBoundaryDisplay.from_estimator(
            clf,
            data,
            multiclass_colors=cmap,
            response_method=response_method,
            plot_method=plot_method,
            ax=ax,
            alpha=0.5,
        )
        ax.scatter(
            data[:, 0],
            data[:, 1],
            c=target.astype(int),
            edgecolors="black",
            cmap=cmap,
        )
        ax.set_title(f"plot_method={plot_method}\nresponse_method={response_method}")