# AUTOGENERATED! DO NOT EDIT! File to edit: 00_inspect.ipynb (unless otherwise specified).

__all__ = ['COLORS', 'Inspector', 'LinRegInspector', 'LogRegInspector']

# Cell
from functools import partial, update_wrapper
from typing import Iterable, Optional, Sequence, Union
import warnings

from IPython.display import HTML
from matplotlib.axes import Axes
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns
from sklearn.base import ClassifierMixin, clone, RegressorMixin
from sklearn import metrics
from sklearn.preprocessing import MinMaxScaler
from sklearn.utils.validation import check_is_fitted
import waterfall_chart

from .explore import plot_correlation
from .tune import calculate_metrics_by_thresh

# Cell
# Meant to be colorblind-friendly
COLORS = {"blue": "#377eb8", "orange": "#ff7f00", "green": "#4daf4a", "pink": "#f781bf"}

# Cell
def Inspector(model, X: pd.DataFrame, y: pd.Series):
    """Model inspector

    Parameters:
    - `model`: Fitted sklearn model
    - `X`: Feature matrix with same shape and column meanings as
    features `model` was trained on
    - `y`: Target series with same length as `X` and same meaning as
    target `model` was trained on
    """
    if hasattr(model, "coef_") and hasattr(model, "intercept_"):
        if isinstance(model, RegressorMixin):
            result = LinRegInspector(model, X, y)
        elif isinstance(model, ClassifierMixin):
            result = LogRegInspector(model, X, y)
    else:
        result = _Inspector(model, X, y)
    return result

# Cell
class _Inspector:
    def __init__(self, model, X, y):
        check_is_fitted(model)

        self.model, self.X, self.y = model, X, y

        self.plot_correlation = partial(
            plot_correlation, pd.concat((self.X, self.y), axis="columns")
        )
        update_wrapper(self.plot_correlation, plot_correlation)

        is_binary = (
            isinstance(self.model, ClassifierMixin) and len(self.y.unique()) == 2
        )
        if is_binary:
            self.calculate_metrics_by_thresh = partial(
                calculate_metrics_by_thresh, self.y, model.predict_proba(self.X)
            )
            update_wrapper(
                self.calculate_metrics_by_thresh, calculate_metrics_by_thresh
            )

        if len(X.columns) == 1:
            if is_binary:
                self.plot = partial(_plot1_bin, self)
                update_wrapper(self.plot, _plot1_bin)
            else:
                self.plot = partial(_plot1, self)
                update_wrapper(self.plot, _plot1)
        if len(X.columns) == 2:
            if isinstance(self.model, ClassifierMixin):
                self.plot = partial(_plot2_clas, self)
                update_wrapper(self.plot, _plot2_clas)
                if is_binary:
                    self.plot3d = partial(_plot3d_bin, self)
                    update_wrapper(self.plot3d, _plot3d_bin)
                else:
                    self.plot3d = partial(_plot3d_multiclass, self)
                    update_wrapper(self.plot3d, _plot3d_multiclass)
            else:
                self.plot = partial(_plot2_regression, self)
                update_wrapper(self.plot, _plot2_regression)
                self.plot3d = partial(_plot3d_regression, self)
                update_wrapper(self.plot3d, _plot3d_regression)

# Cell
class _LinearInspector(_Inspector):
    def __init__(self, model, X, y):
        super().__init__(model, X, y)
        if isinstance(self.model, ClassifierMixin) and len(self.y.unique()) == 2:
            self.plot_waterfall = partial(_plot_waterfall_bin, self)
            update_wrapper(self.plot_waterfall, _plot_waterfall_bin)

    def show_equation(*args, **kwargs):
        raise NotImplementedError()

    def plot_coefs_vs_hparam(*args, **kwargs):
        raise NotImplementedError()

# Cell
class LinRegInspector(_LinearInspector):
    """Linear regression model inspector"""

    def plot_coefs_vs_hparam(self, hparam: str, vals: Sequence[float]):
        """Plot coefficient values against a hyperparameter

        Parameters:
        - `hparam`: Name of hyperparameter; must be an attribute of
        `self.model`
        - `vals`: Values of that hyperparameter to use
        """
        current_val = getattr(self.model, hparam)
        model = clone(self.model)
        setattr(model, hparam, vals[-1])
        model.fit(self.X, self.y)
        column_order = model.coef_.argsort()[::-1]
        X = self.X.iloc[:, column_order]

        coefs = []
        for val in vals:
            setattr(model, hparam, val)
            coefs.append(model.fit(X, self.y).coef_)

        fig, ax = plt.subplots()
        ax.plot(vals, coefs)
        ax.axvline(current_val, c="k", label="current value")
        ax.set(xlabel=hparam, ylabel="coefficient value")
        ax.legend(X.columns, bbox_to_anchor=(1.05, 1.0), loc="upper left")
        return ax

    def plot_waterfall(
        self,
        item: Union[pd.Series, np.array],
        bar_num_formatter: str = ".1f",
        tick_num_formatter: str = ".2f",
        waterfall_kwargs: Optional[dict] = None,
    ):
        """Make a waterfall chart showing how each feature contributes
        to the prediction for the input item.

        Parameters:
        - `item`: Input item, with the same shape and value meanings as
        a single row from `self.X`
        - `bar_num_formatter`: Bar label format specifier
        - `tick_num_formatter`: Tick label format specifier
        - `waterfall_kwargs`: kwargs to pass to `waterfall_chart.plot`
        """
        if waterfall_kwargs is None:
            waterfall_kwargs = {
                "sorted_value": True,
                "threshold": 0.01,
                "blue_color": COLORS["blue"],
                "green_color": COLORS["green"],
                "red_color": COLORS["orange"],
            }
        index = ["int"] + [
            f"{name}: {val:{tick_num_formatter}}"
            for name, val in zip(self.X.columns, item)
        ]
        vals = [self.model.intercept_] + list(self.model.coef_ * item)
        waterfall_chart.plot(
            index=index,
            data=vals,
            x_lab="Feature name and value",
            y_lab="Contribution to prediction",
            formatting=f"{{:,{bar_num_formatter}}}",
            net_label=self.y.name,
            **waterfall_kwargs,
        )
        return plt.gca()

    def show_equation(
        self,
        intercept_formatter: str = ".2f",
        coef_formatter: str = ".2f",
    ):
        """Show linear model equation

        Parameters:
        - `intercept_formatter`: Intercept format specifier
        - `coef_formatter`: Intercept format specifier
        """
        return HTML(
            _generate_linear_model_html(
                intercept=self.model.intercept_,
                coefs=self.model.coef_,
                feature_names=self.X.columns,
                target_name=self.y.name,
                intercept_formatter=intercept_formatter,
                coef_formatter=coef_formatter,
            )
        )

# Cell
class LogRegInspector(_LinearInspector):
    """Logistic regression model inspector"""

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        is_binary = len(self.y.unique()) == 2
        if is_binary:
            self.show_equation = partial(_show_equation_bin, self)
            update_wrapper(self.show_equation, _show_equation_bin)
        else:
            self.show_equation = partial(_show_equation_multiclass, self)
            update_wrapper(self.show_equation, _show_equation_multiclass)

    def plot_coefs_vs_hparam(self, hparam: str, vals: Sequence[float]) -> np.array:
        """Plot coefficient values against a hyperparameter

        Parameters:
        - `hparam`: Name of hyperparameter; must be an attribute of
        `self.model`
        - `vals`: Values of that hyperparameter to use
        """
        current_val = getattr(self.model, hparam)
        model = clone(self.model)
        setattr(model, hparam, vals[-1])
        column_order = model.fit(self.X, self.y).coef_[0].argsort()[::-1]
        X = self.X.iloc[:, column_order]

        coef_arrays = []
        for val in vals:
            setattr(model, hparam, val)
            coef_arrays.append(model.fit(self.X, self.y).coef_)

        num_target_vals = len(set(self.y))
        if num_target_vals == 2:
            fig, ax = plt.subplots()
            ax.plot(vals, [coefs[0] for coefs in coef_arrays])
            axes = np.array(ax)[None]
        else:
            fig, axes = plt.subplots(
                num_target_vals, 1, sharex=True, sharey=True, constrained_layout=True
            )
            for target_val_num in range(num_target_vals):
                axes[target_val_num].plot(
                    vals, [coefs[target_val_num] for coefs in coef_arrays]
                )
                axes[target_val_num].set_title(
                    f"y={sorted(set(self.y))[target_val_num]}"
                )
        axes[0].set(xlabel=hparam, ylabel="Coefficient Value")
        for ax in axes:
            ax.axvline(current_val, c="k", label="current value")
        axes[0].legend(X.columns, bbox_to_anchor=(1.05, 1.0), loc="upper left")
        return axes

# Cell
def _show_equation_multiclass(
    inspector,
    intercept_formatter: str = ".2f",
    coef_formatter: str = ".2f",
):
    """Show multiclass logistic model equation

    Parameters:
    - `intercept_formatter`: Intercept format specifier
    - `coef_formatter`: Intercept format specifier
    """
    model_string = ""
    for target_name, coefs, intercept in zip(
        self.y.unique(), self.model.coef_, self.model.intercept_
    ):
        model_string += f"""
                <p>
                    {_generate_linear_model_html(
                            intercept=intercept,
                            coefs=coefs,
                            feature_names=self.X.columns,
                            target_name=f"log-odds({self.y.name} = {target_name})",
                            intercept_formatter=intercept_formatter,
                            coef_formatter=coef_formatter,
                        )
                    }
                </p>
            """
    return HTML(model_string)

# Cell
def _show_equation_bin(
    inspector,
    intercept_formatter: str = ".2f",
    coef_formatter: str = ".2f",
):
    """Show binary logistic model equation

    Parameters:
    - `intercept_formatter`: Intercept format specifier
    - `coef_formatter`: Intercept format specifier
    """
    return HTML(
        _generate_linear_model_html(
            intercept=inspector.model.intercept_[0],
            coefs=inspector.model.coef_[0],
            feature_names=inspector.X.columns,
            target_name=f"log-odds({inspector.y.name})",
            intercept_formatter=intercept_formatter,
            coef_formatter=coef_formatter,
        )
    )

# Cell
def _plot1(
    inspector: Inspector,
    plot_data: bool = True,
    line_kwargs: Optional[dict] = None,
    scatter_kwargs: Optional[dict] = None,
    ax: Optional[Axes] = None,
) -> Axes:
    """Plot predictions from a regression or multiclass model with a
    single input as a line

    Parameters:
    - `inspector`: `Inspector` object
    - `plot_data`: Make a scatter plot of the data
    - `line_kwargs`: kwargs to pass to `ax.plot` for plotting
    predictions
    - `scatter_kwargs`: kwargs to pass to `ax.scatter` for plotting data
    - `ax`: Matplotlib `Axes` object. Plot will be added to this object
    if provided; otherwise a new `Axes` object will be generated.
    """

    def _plot_preds(ax):
        X_sorted = inspector.X.sort_values(inspector.X.columns[0])
        ax.plot(
            X_sorted.iloc[:, 0],
            inspector.model.predict(X_sorted),
            label="predictions",
            **line_kwargs,
        )
        return ax

    if line_kwargs is None:
        line_kwargs = {}
    if scatter_kwargs is None:
        scatter_kwargs = {"c": "k", "alpha": 0.4}
    if ax is None:
        _, ax = plt.subplots()
    if plot_data:
        ax.scatter(inspector.X.iloc[:, 0], inspector.y, **scatter_kwargs)
    ax = _plot_preds(ax)
    ax.set(xlabel=inspector.X.columns[0], ylabel=inspector.y.name)
    ax.legend()
    return ax

# Cell
def _plot1_bin(
    inspector: Inspector,
    thresh: Optional[float] = 0.5,
    plot_data: bool = True,
    prob_line_kwargs: Optional[dict] = None,
    thresh_line_kwargs: Optional[dict] = None,
    scatter_kwargs: Optional[dict] = None,
    ax: Optional[Axes] = None,
) -> Axes:
    """Plot predictions from a binary classification model that provides
    probabilities and has a single input

    Parameters:
    - `inspector`: `Inspector` object
    - `thresh`: Threshold probability
    - `plot_data`: Make a scatter plot of the data
    - `prob_line_kwargs`: kwargs to pass to `ax.plot` for plotting
    model probabilities
    - `thresh_line_kwargs`: kwargs to pass to `ax.plot` for plotting
    threshold
    - `scatter_kwargs`: kwargs to pass to `ax.scatter` for plotting data
    - `ax`: Matplotlib `Axes` object. Plot will be added to this object
    if provided; otherwise a new `Axes` object will be generated.
    """

    def _plot_probs(ax):
        num_points = 100
        X = np.linspace(inspector.X.min(), inspector.X.max(), num_points)
        ax.plot(
            X,
            inspector.model.predict_proba(X)[:, 1],
            label="probability",
            **prob_line_kwargs,
        )
        return ax

    if ax is None:
        _, ax = plt.subplots()
    if prob_line_kwargs is None:
        prob_line_kwargs = {}
    if thresh_line_kwargs is None:
        thresh_line_kwargs = {}
    if scatter_kwargs is None:
        scatter_kwargs = {"c": "k", "alpha": 0.4}

    if plot_data:
        ax.scatter(inspector.X.iloc[:, 0], inspector.y, **scatter_kwargs)
    ax = _plot_probs(ax)
    if thresh:
        ax.plot(
            inspector.X.iloc[:, 0],
            thresh * np.ones(inspector.X.shape),
            **thresh_line_kwargs,
        )
    ax.set(xlabel=inspector.X.columns[0], ylabel=inspector.y.name)
    ax.legend()
    return ax

# Cell
def _plot2_regression(
    inspector: Inspector,
    plot_data: bool = True,
    heatmap_kwargs: Optional[dict] = None,
    scatter_kwargs: Optional[dict] = None,
    tick_formatter: Optional[str] = ".2f",
    ax=None,
):
    """Plot predictions from a model with two inputs as a heatmap.

    Parameters:
    - `inspector`: `Inspector` object
    - `plot_data`: Make a scatter plot of the data
    - `heatmap_kwargs`: kwargs to pass to `sns.heatmap` for plotting
    predictions
    - `scatter_kwargs`: kwargs to pass to `ax.scatter` for plotting data
    - `tick_formatter`: Tick label format specifier
    - `ax`: Matplotlib `Axes` object. Plot will be added to this object
    if provided; otherwise a new `Axes` object will be generated.
    """

    def _plot_preds(ax, **heatmap_kwargs):
        x_grid = np.linspace(
            inspector.X.iloc[:, 0].min(), inspector.X.iloc[:, 0].max(), 100
        )
        y_grid = np.linspace(
            inspector.X.iloc[:, 1].max(), inspector.X.iloc[:, 1].min(), 100
        )

        preds = inspector.model.predict(
            np.transpose([np.tile(x_grid, len(y_grid)), np.repeat(y_grid, len(x_grid))])
        ).reshape(len(y_grid), len(x_grid))
        preds = pd.DataFrame(preds, columns=x_grid, index=y_grid)
        return sns.heatmap(
            preds,
            vmin=inspector.y.min(),
            vmax=inspector.y.max(),
            ax=ax,
            **heatmap_kwargs,
        )

    if ax is None:
        _, ax = plt.subplots()
    if heatmap_kwargs is None:
        heatmap_kwargs = {"cmap": "viridis"}
    if scatter_kwargs is None:
        scatter_kwargs = {"cmap": "viridis", "edgecolor": "k", "zorder": 999}

    if plot_data:
        ax = _plot_data_2d(X=inspector.X, y=inspector.y, ax=ax, **scatter_kwargs)
    ax = _plot_preds(ax=ax, **heatmap_kwargs)
    if tick_formatter is not None:
        _format_ticks(ax=ax, formatter=tick_formatter)
    ax.set(
        xlabel=inspector.X.columns[0],
        ylabel=inspector.X.columns[1],
        title=inspector.y.name,
    )
    return ax

# Cell
def _plot2_clas(
    inspector: Inspector,
    plot_data: bool = True,
    heatmap_kwargs: Optional[dict] = None,
    scatter_kwargs: Optional[dict] = None,
    tick_formatter: Optional[str] = ".2f",
    ax=None,
):
    """Plot data and predictions for classification model with two features

    Parameters:
    - `inspector`: `Inspector` object
    - `plot_data`: Make a scatter plot of the data
    - `heatmap_kwargs`: kwargs to pass to `sns.heatmap` for plotting
    predictions
    - `scatter_kwargs`: kwargs to pass to `ax.scatter` for plotting data
    - `tick_formatter`: Tick label format specifier
    - `ax`: Matplotlib `Axes` object. Plot will be added to this object
    if provided; otherwise a new `Axes` object will be generated.
    """

    def _plot_preds(y_vals, label_to_num, ax, **scatter_kwargs):
        num_points = 100
        x_grid = np.linspace(
            inspector.X.iloc[:, 0].min(), inspector.X.iloc[:, 0].max(), num_points
        )
        y_grid = np.linspace(
            inspector.X.iloc[:, 1].max(), inspector.X.iloc[:, 1].min(), num_points
        )

        preds = inspector.model.predict(
            np.transpose([np.tile(x_grid, len(y_grid)), np.repeat(y_grid, len(x_grid))])
        ).reshape(len(y_grid), len(x_grid))
        preds = pd.DataFrame(preds, columns=x_grid, index=y_grid)
        for col in preds:
            preds.loc[:, col] = preds.loc[:, col].map(label_to_num)
        ax = sns.heatmap(preds.astype(int), **heatmap_kwargs)
        return ax

    def _set_colorbar(y_vals, ax):
        colorbar = ax.collections[0].colorbar
        r = colorbar.vmax - colorbar.vmin
        colorbar.set_ticks(
            [colorbar.vmin + r / len(y_vals) * (0.5 + i) for i in range(len(y_vals))]
        )
        colorbar.set_ticklabels(y_vals)
        return colorbar

    def _wash_out(ax):
        rectangle = plt.Rectangle((0, 0), 100, 100, fc="w", alpha=0.5)
        ax.add_patch(rectangle)
        return ax

    if ax is None:
        _, ax = plt.subplots()
    if heatmap_kwargs is None:
        heatmap_kwargs = {}
    if scatter_kwargs is None:
        scatter_kwargs = {"edgecolor": "k", "zorder": 999}

    y_vals = inspector.y.unique()
    label_to_num = {label: num for label, num in zip(y_vals, range(len(y_vals)))}
    if heatmap_kwargs.get("cmap") is None:
        heatmap_kwargs["cmap"] = sns.color_palette(None, len(y_vals))

    ax = _plot_preds(y_vals, label_to_num, ax=ax, **heatmap_kwargs)
    ax = _wash_out(ax)
    colorbar = _set_colorbar(y_vals=y_vals, ax=ax)

    if plot_data:
        if scatter_kwargs.get("cmap") is None:
            scatter_kwargs["cmap"] = colorbar.cmap
        ax = _plot_data_2d(
            X=inspector.X, y=inspector.y.map(label_to_num), ax=ax, **scatter_kwargs
        )
    _format_ticks(ax=ax, formatter=tick_formatter)
    return ax

# Cell
def _plot_data_2d(X, y, ax, **scatter_kwargs):
    X_normalized = MinMaxScaler().fit_transform(X) * 99
    ax.scatter(
        X_normalized[:, 0] + 0.5,
        X_normalized[:, 1].max() - X_normalized[:, 1] + 0.5,
        c=y,
        **scatter_kwargs,
    )
    ax.set(xlabel=X.columns[0], ylabel=X.columns[1])
    return ax

# Cell
def _format_ticks(ax, formatter):
    labels = [item.get_text() for item in ax.get_xticklabels()]
    ax.set_xticklabels([f"{float(label):{formatter}}" for label in labels])

    labels = [item.get_text() for item in ax.get_yticklabels()]
    ax.set_yticklabels([f"{float(label):{formatter}}" for label in labels])

# Cell
def _plot3d_regression(
    inspector: Inspector,
    plot_data: bool = True,
    surf_kwargs: Optional[dict] = None,
    scatter_kwargs: Optional[dict] = None,
    ax=None,
):
    """Plot predictions from a model with two inputs as a surface.

    Best viewed with a tool such as https://github.com/matplotlib/ipympl
    that supports rotating the output

    Parameters:
    - `inspector`: `Inspector` object
    - `plot_data`: Make a scatter plot of the data
    - `surf_kwargs`: kwargs to pass to `ax.plot_surface` for plotting
    predictions
    - `scatter_kwargs`: kwargs to pass to `ax.scatter` for plotting data
    - `ax`: Matplotlib `Axes` object. Plot will be added to this object
    if provided; otherwise a new `Axes` object will be generated.
    """

    def _plot_preds(ax):

        x0_grid, x1_grid = _create_2d_grid(inspector.X)
        ax.plot_surface(
            x0_grid,
            x1_grid,
            _get_2d_grid_preds(inspector.model, x0_grid, x1_grid),
            rstride=1,
            cstride=1,
            vmin=inspector.y.min(),
            vmax=inspector.y.max(),
            **surf_kwargs,
        )
        return ax

    if ax is None:
        fig = plt.figure()
        ax = fig.add_subplot(111, projection="3d")
    if surf_kwargs is None:
        surf_kwargs = {"alpha": 0.4, "cmap": "viridis"}
    if scatter_kwargs is None:
        scatter_kwargs = {"cmap": "viridis"}
    if plot_data:
        ax.scatter(
            inspector.X.iloc[:, 0],
            inspector.X.iloc[:, 1],
            inspector.y,
            c=inspector.y,
            **scatter_kwargs,
        )
    ax = _plot_preds(ax)
    ax.set(
        xlabel=inspector.X.columns[0],
        ylabel=inspector.X.columns[1],
        zlabel=inspector.y.name,
    )
    return ax

# Cell
def _plot3d_multiclass(
    inspector: Inspector,
    plot_data: bool = True,
    surf_kwargs: Optional[dict] = None,
    scatter_kwargs: Optional[dict] = None,
    ax=None,
):
    """Plot predictions from a model with two inputs as a surface.

    Best viewed with a tool such as https://github.com/matplotlib/ipympl
    that supports rotating the output

    Parameters:
    - `inspector`: `Inspector` object
    - `plot_data`: Make a scatter plot of the data
    - `surf_kwargs`: kwargs to pass to `ax.plot_surface` for plotting
    predictions
    - `scatter_kwargs`: kwargs to pass to `ax.scatter` for plotting data
    - `ax`: Matplotlib `Axes` object. Plot will be added to this object
    if provided; otherwise a new `Axes` object will be generated.
    """
    if ax is None:
        fig = plt.figure()
        ax = fig.add_subplot(111, projection="3d")
    if surf_kwargs is None:
        surf_kwargs = {"alpha": 0.4, "cmap": "viridis"}
    if scatter_kwargs is None:
        scatter_kwargs = {}
    y_vals = inspector.y.unique()
    label_to_num = {label: num for label, num in zip(y_vals, range(len(y_vals)))}
    y_int = inspector.y.map(label_to_num)

    y_pred_int = pd.Series(inspector.model.predict(inspector.X)).map(label_to_num)
    x0_grid, x1_grid = _create_2d_grid(inspector.X, num_points=20)
    grid_preds = pd.DataFrame(
        _get_2d_grid_preds(inspector.model, x0_grid, x1_grid)
    ).applymap(lambda x: label_to_num[x])

    for val in y_int.unique():
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            ax.plot_surface(
                x0_grid,
                x1_grid,
                grid_preds[grid_preds == val],
                rstride=1,
                cstride=1,
                alpha=0.3,
            )
        if plot_data:
            ax.scatter(
                inspector.X.iloc[:, 0].loc[y_int == val],
                inspector.X.iloc[:, 1].loc[y_int == val],
                y_pred_int.loc[y_int == val],
                **scatter_kwargs,
            )

    return ax

# Cell
def _plot3d_bin(
    inspector: Inspector,
    thresh=0.5,
    plot_prob: bool = True,
    plot_thresh: bool = True,
    plot_data: bool = True,
    prob_surf_kwargs: Optional[dict] = None,
    thresh_surf_kwargs: Optional[dict] = None,
    scatter_kwargs: Optional[dict] = None,
    ax=None,
):
    """Plot data and predictions for binary classification model with 2
    features in 3D

    Best viewed with a tool such as https://github.com/matplotlib/ipympl
    that supports rotating the output

    Parameters:
    - `inspector`: `Inspector` object
    - `thresh`: Probability threshold for counting a prediction as
    positive
    - `plot_prob`: Whether to plot the model probabilities
    - `plot_thresh`: Whether to plot a classification threshold
    - `plot_data`: Whether to plot the data
    - `prob_surf_kwargs`: kwargs to pass to the model probability
    surface
    - `thresh_surf_kwargs`: kwargs to pass to the threshold surface
    - `scatter_kwargs`: kwargs to pass to the scatter plot of the data
    - `ax`: Matplotlib `Axes` object. Plot will be added to this object
    if provided; otherwise a new `Axes` object will be generated.
    """

    def _get_grid_probs():
        return inspector.model.predict_proba(
            np.hstack((x0_grid.reshape(-1, 1), x1_grid.reshape(-1, 1)))
        )[:, 1].reshape(x0_grid.shape)

    if ax is None:
        fig = plt.figure()
        ax = fig.add_subplot(111, projection="3d")
    if prob_surf_kwargs is None:
        prob_surf_kwargs = {"alpha": 0.4, "cmap": "viridis"}
    if thresh_surf_kwargs is None:
        thresh_surf_kwargs = {"alpha": 0.4, "color": "k"}
    if scatter_kwargs is None:
        scatter_kwargs = {}

    x0_grid, x1_grid = _create_2d_grid(inspector.X)

    if plot_prob:
        ax.plot_surface(
            x0_grid,
            x1_grid,
            _get_grid_probs(),
            rstride=1,
            cstride=1,
            **prob_surf_kwargs,
        )
    if plot_data:
        y_pred = inspector.model.predict_proba(inspector.X)[:, 1] > thresh
        ax.scatter(
            inspector.X.loc[y_pred == inspector.y].iloc[:, 0],
            inspector.X.loc[y_pred == inspector.y].iloc[:, 1],
            inspector.y.loc[y_pred == inspector.y],
            **scatter_kwargs,
            label="correct",
        )
        ax.scatter(
            inspector.X.loc[y_pred != inspector.y].iloc[:, 0],
            inspector.X.loc[y_pred != inspector.y].iloc[:, 1],
            inspector.y.loc[y_pred != inspector.y],
            **scatter_kwargs,
            label="incorrect",
        )
        ax.legend()
    if plot_thresh:
        ax.plot_surface(
            x0_grid,
            x1_grid,
            thresh * np.ones((len(x0_grid), len(x1_grid))),
            rstride=1,
            cstride=1,
            **thresh_surf_kwargs,
        )
    ax.set(
        xlabel=inspector.X.columns[0],
        ylabel=inspector.X.columns[1],
        zlabel=f"{inspector.y.name} prediction",
    )
    return ax

# Cell
def _create_2d_grid(X, num_points=20):
    x0_grid = np.linspace(X.iloc[:, 0].min(), X.iloc[:, 0].max(), num_points)
    x1_grid = np.linspace(X.iloc[:, 1].min(), X.iloc[:, 1].max(), num_points)
    return np.meshgrid(x0_grid, x1_grid)

# Cell
def _get_2d_grid_preds(model, x0_grid, x1_grid):
    return model.predict(
        np.hstack((x0_grid.reshape(-1, 1), x1_grid.reshape(-1, 1)))
    ).reshape(x0_grid.shape)

# Cell
def _plot_waterfall_bin(
    interp,
    item: Union[pd.Series, np.array],
    bar_num_formatter: str = ".1f",
    tick_num_formatter: str = ".2f",
    waterfall_kwargs: Optional[dict] = None,
):
    """Make a waterfall chart showing how each feature contributes
    to the prediction for the input item for a binary classification
    model.

    Parameters:
    - `item`: Input item, with the same shape and value meanings as
    a single row from `interp.X`
    - `bar_num_formatter`: Bar label format specifier
    - `tick_num_formatter`: Tick label format specifier
    - ``waterfall_kwargs`: kwargs to pass to `waterfall_chart.plot`
    """
    if waterfall_kwargs is None:
        waterfall_kwargs = {
            "sorted_value": True,
            "threshold": 0.01,
            "blue_color": COLORS["blue"],
            "green_color": COLORS["green"],
            "red_color": COLORS["orange"],
        }

    index = ["int"] + [
        f"{name}: {val:{tick_num_formatter}}"
        for name, val in zip(interp.X.columns, item)
    ]
    vals = [interp.model.intercept_[0]] + list(interp.model.coef_[0] * item)
    waterfall_chart.plot(
        index=index,
        data=vals,
        x_lab="Feature name and value",
        y_lab="Contribution to log-odds",
        formatting=f"{{:,{bar_num_formatter}}}",
        net_label=interp.y.name,
        **waterfall_kwargs,
    )
    return plt.gca()

# Cell
def _generate_linear_model_html(
    intercept: float,
    coefs: Sequence[float],
    feature_names: Iterable[str],
    target_name: str,
    intercept_formatter: str = ".2f",
    coef_formatter: str = ".2f",
):
    if len(coefs) != len(feature_names):
        raise ValueError("len(coefs) != len(feature_cols)")
    model_string = f"""
        <span style='color:{COLORS["pink"]}'>{target_name}</span>
        = <span style='color:{COLORS["orange"]}'>{intercept:{intercept_formatter}}</span>
    """
    for coef, feature_col in zip(coefs, feature_names):
        model_string += f"""
            <span style='color:{COLORS["green"]}'>{"+" if coef >= 0 else "-"} {abs(coef):{coef_formatter}}</span>
            * <span style='color:{COLORS["blue"]}'>{feature_col}</span>
        """
    return model_string