import numpy as np
import warnings
try:
    import matplotlib.pyplot as plt
    import matplotlib
except ImportError:
    warnings.warn("matplotlib could not be loaded!")
    pass
from .labels import labels
from .general import safe_isinstance, format_value
from . import colors

# 221116 Added argument (y_reverse) to function waterfall that reverses the y axis when set to True
# (y_reverse = True: Expected value at top of y axis, Predicted value at bottom of y axis)

# TODO: If we make a JS version of this plot then we could let users click on a bar and then see the dependence
# plot that is associated with that feature get overlayed on the plot...it would quickly allow users to answer
# why a feature is pushing down or up. Perhaps the best way to do this would be with an ICE plot hanging off
# of the bar...
def waterfall(shap_values, max_display=10, show=True, y_reverse=False, rank_absolute=True, raw_ascending=True):
    """ Plots an explantion of a single prediction as a waterfall plot.
    The SHAP value of a feature represents the impact of the evidence provided by that feature on the model's
    output. The waterfall plot is designed to visually display how the SHAP values (evidence) of each feature
    move the model output from our prior expectation under the background data distribution, to the final model
    prediction given the evidence of all the features. Features are sorted by the magnitude of their SHAP values
    with the smallest magnitude features grouped together at the bottom of the plot when the number of features
    in the models exceeds the max_display parameter.
    Parameters
    ----------
    shap_values : Explanation
        A one-dimensional Explanation object that contains the feature values and SHAP values to plot.
    max_display : str
        The maximum number of features to plot.
    show : bool
        Whether matplotlib.pyplot.show() is called before returning. Setting this to False allows the plot
        to be customized further after it has been created.
    y_reverse : bool
        Determines placement of expected value & predicted value on y axis. Setting this to True reverses the 
        y axis and puts expected value at top of y axis, and predicted value at bottom of y axis.
    rank_absolute : bool
        If True then rank values by their absolute values.
    raw_ascending : bool
        If rank_absolute == False then look to this setting, else redundant (as order is on absolute values).
        Determines order of ranked features if not absolute. If True then from expected value put features in 
        order of most negative to least negative to smallest positive to most positive. If False then from expected 
        value put features in order of most positive to smallest positive to least negative to most negative.
    """
    # Turn off interactive plot
    if show is False:
        plt.ioff()

    base_values = shap_values.base_values
    features = shap_values.display_data if shap_values.display_data is not None else shap_values.data
    feature_names = shap_values.feature_names
    lower_bounds = getattr(shap_values, "lower_bounds", None)
    upper_bounds = getattr(shap_values, "upper_bounds", None)
    values = shap_values.values

    # make sure we only have a single output to explain
    if (type(base_values) == np.ndarray and len(base_values) > 0) or type(base_values) == list:
        raise Exception("waterfall_plot requires a scalar base_values of the model output as the first "
                        "parameter, but you have passed an array as the first parameter! "
                        "Try shap.waterfall_plot(explainer.base_values[0], values[0], X[0]) or "
                        "for multi-output models try "
                        "shap.waterfall_plot(explainer.base_values[0], values[0][0], X[0]).")

    # make sure we only have a single explanation to plot
    if len(values.shape) == 2:
        raise Exception(
            "The waterfall_plot can currently only plot a single explanation but a matrix of explanations was passed!")

    # unwrap pandas series
    if safe_isinstance(features, "pandas.core.series.Series"):
        if feature_names is None:
            feature_names = list(features.index)
        features = features.values

    # fallback feature names
    if feature_names is None:
        feature_names = np.array([labels['FEATURE'] % str(i) for i in range(len(values))])

    # init variables we use for tracking the plot locations
    num_features = min(max_display, len(values))
    row_height = 0.5
    rng = range(num_features - 1, -1, -1)
    
    # Rank features by absolute values to choose which are going to be included individually in the 
    #   waterfall plot, the rest will be represented in one summed bar.
    order = np.argsort(-np.abs(values))

    pos_lefts = []
    pos_inds = []
    pos_widths = []
    pos_low = []
    pos_high = []
    neg_lefts = []
    neg_inds = []
    neg_widths = []
    neg_low = []
    neg_high = []
    loc = base_values + values.sum()
    yticklabels = ["" for i in range(num_features + 1)]

    # size the plot based on how many features we are plotting
    plt.gcf().set_size_inches(8, num_features * row_height + 1.5)

    # see how many individual (vs. grouped at the end) features we are plotting
    if num_features == len(values):
        num_individual = num_features
    else:
        num_individual = num_features - 1
        
    if not rank_absolute:
        # The plotting order of just the features that are going to be shown indiviually to use raw values
        if raw_ascending:
            # in ascending order (most negative to minimal change to most positive)
            values_to_sort = [values[order[i]] for i in range(num_individual)]
            arg_to_sort = order[0:num_individual]
            arg_values_to_sort = np.argsort(values_to_sort)
            # range(num_individual-1,-1,-1) makes to loop go backwards from n to 0.
            order_arg_to_sort = [arg_to_sort[arg_values_to_sort[i]] for i in range(num_individual-1,-1,-1)]
            # Replace the location values for just those features that are going to be shown individually.
            order[0:num_individual] = order_arg_to_sort[0:num_individual]
        else:
            # in descending order (most positive to minimal change to most negative)
            # sort the -1*values
            values_to_sort = [-values[order[i]] for i in range(num_individual)]
            arg_to_sort = order[0:num_individual]
            arg_values_to_sort = np.argsort(values_to_sort)
            # range(num_individual-1,-1,-1) makes to loop go backwards from n to 0.
            order_arg_to_sort = [arg_to_sort[arg_values_to_sort[i]] for i in range(num_individual-1,-1,-1)]
            # Replace the location values for just those features that are going to be shown individually.
            order[0:num_individual] = order_arg_to_sort[0:num_individual]

    # compute the locations of the individual features and plot the dashed connecting lines
    for i in range(num_individual):
        sval = values[order[i]]
        loc -= sval
        if sval >= 0:
            pos_inds.append(rng[i])
            pos_widths.append(sval)
            if lower_bounds is not None:
                pos_low.append(lower_bounds[order[i]])
                pos_high.append(upper_bounds[order[i]])
            pos_lefts.append(loc)
        else:
            neg_inds.append(rng[i])
            neg_widths.append(sval)
            if lower_bounds is not None:
                neg_low.append(lower_bounds[order[i]])
                neg_high.append(upper_bounds[order[i]])
            neg_lefts.append(loc)
        if num_individual != num_features or i + 4 < num_individual:
            plt.plot([loc, loc], [rng[i] - 1 - 0.4, rng[i] + 0.4],
                     color="#bbbbbb", linestyle="--", linewidth=0.5, zorder=-1)
        if features is None:
            yticklabels[rng[i]] = feature_names[order[i]]
        else:
            if np.issubdtype(type(features[order[i]]), np.number):
                yticklabels[rng[i]] = format_value(float(features[order[i]]), "%0.03f") + " = " + feature_names[order[i]]
            else:
                yticklabels[rng[i]] = features[order[i]] + " = " + feature_names[order[i]]

    # add a last grouped feature to represent the impact of all the features we didn't show
    if num_features < len(values):
        yticklabels[0] = "%d other features" % (len(values) - num_features + 1)
        remaining_impact = base_values - loc
        if remaining_impact < 0:
            pos_inds.append(0)
            pos_widths.append(-remaining_impact)
            pos_lefts.append(loc + remaining_impact)
            c = colors.red_rgb
        else:
            neg_inds.append(0)
            neg_widths.append(-remaining_impact)
            neg_lefts.append(loc + remaining_impact)
            c = colors.blue_rgb

    points = pos_lefts + list(np.array(pos_lefts) + np.array(pos_widths)) + neg_lefts + \
        list(np.array(neg_lefts) + np.array(neg_widths))
    dataw = np.max(points) - np.min(points)

    # draw invisible bars just for sizing the axes
    label_padding = np.array([0.1*dataw if w < 1 else 0 for w in pos_widths])
    plt.barh(pos_inds, np.array(pos_widths) + label_padding + 0.02*dataw,
             left=np.array(pos_lefts) - 0.01*dataw, color=colors.red_rgb, alpha=0)
    label_padding = np.array([-0.1*dataw if -w < 1 else 0 for w in neg_widths])
    plt.barh(neg_inds, np.array(neg_widths) + label_padding - 0.02*dataw,
             left=np.array(neg_lefts) + 0.01*dataw, color=colors.blue_rgb, alpha=0)

    # define variable we need for plotting the arrows
    head_length = 0.08
    bar_width = 0.8
    xlen = plt.xlim()[1] - plt.xlim()[0]
    fig = plt.gcf()
    ax = plt.gca()
    xticks = ax.get_xticks()
    bbox = ax.get_window_extent().transformed(fig.dpi_scale_trans.inverted())
    width, height = bbox.width, bbox.height
    bbox_to_xscale = xlen/width
    hl_scaled = bbox_to_xscale * head_length
    renderer = fig.canvas.get_renderer()

    # draw the positive arrows
    for i in range(len(pos_inds)):
        dist = pos_widths[i]
        arrow_obj = plt.arrow(
            pos_lefts[i], pos_inds[i], max(dist-hl_scaled, 0.000001), 0,
            head_length=min(dist, hl_scaled),
            color=colors.red_rgb, width=bar_width,
            head_width=bar_width
        )

        if pos_low is not None and i < len(pos_low):
            plt.errorbar(
                pos_lefts[i] + pos_widths[i], pos_inds[i],
                xerr=np.array([[pos_widths[i] - pos_low[i]], [pos_high[i] - pos_widths[i]]]),
                ecolor=colors.light_red_rgb
            )

        txt_obj = plt.text(
            pos_lefts[i] + 0.5*dist, pos_inds[i], format_value(pos_widths[i], '%+0.02f'),
            horizontalalignment='center', verticalalignment='center', color="white",
            fontsize=12
        )
        text_bbox = txt_obj.get_window_extent(renderer=renderer)
        arrow_bbox = arrow_obj.get_window_extent(renderer=renderer)

        # if the text overflows the arrow then draw it after the arrow
        if text_bbox.width > arrow_bbox.width:
            txt_obj.remove()

            txt_obj = plt.text(
                pos_lefts[i] + (5/72)*bbox_to_xscale + dist, pos_inds[i], format_value(pos_widths[i], '%+0.02f'),
                horizontalalignment='left', verticalalignment='center', color=colors.red_rgb,
                fontsize=12
            )

    # draw the negative arrows
    for i in range(len(neg_inds)):
        dist = neg_widths[i]

        arrow_obj = plt.arrow(
            neg_lefts[i], neg_inds[i], -max(-dist-hl_scaled, 0.000001), 0,
            head_length=min(-dist, hl_scaled),
            color=colors.blue_rgb, width=bar_width,
            head_width=bar_width
        )

        if neg_low is not None and i < len(neg_low):
            plt.errorbar(
                neg_lefts[i] + neg_widths[i], neg_inds[i],
                xerr=np.array([[neg_widths[i] - neg_low[i]], [neg_high[i] - neg_widths[i]]]),
                ecolor=colors.light_blue_rgb
            )

        txt_obj = plt.text(
            neg_lefts[i] + 0.5*dist, neg_inds[i], format_value(neg_widths[i], '%+0.02f'),
            horizontalalignment='center', verticalalignment='center', color="white",
            fontsize=12
        )
        text_bbox = txt_obj.get_window_extent(renderer=renderer)
        arrow_bbox = arrow_obj.get_window_extent(renderer=renderer)

        # if the text overflows the arrow then draw it after the arrow
        if text_bbox.width > arrow_bbox.width:
            txt_obj.remove()

            txt_obj = plt.text(
                neg_lefts[i] - (5/72)*bbox_to_xscale + dist, neg_inds[i], format_value(neg_widths[i], '%+0.02f'),
                horizontalalignment='right', verticalalignment='center', color=colors.blue_rgb,
                fontsize=12
            )

    # draw the y-ticks twice, once in gray and then again with just the feature names in black
    # The 1e-8 is so matplotlib 3.3 doesn't try and collapse the ticks
    ytick_pos = list(range(num_features)) + list(np.arange(num_features)+1e-8)
    plt.yticks(ytick_pos, yticklabels[:-1] + [l.split('=')[-1] for l in yticklabels[:-1]], fontsize=13)

    # put horizontal lines for each feature row
    for i in range(num_features):
        plt.axhline(i, color="#cccccc", lw=0.5, dashes=(1, 5), zorder=-1)

    if y_reverse:
        ex_loc = 1
        fx_loc = 1/num_features
    else:
        ex_loc = 1/num_features
        fx_loc = 1
    # mark the prior expected value and the model prediction with a vertical dashed grey line
    # depending on if flipping the y axis, as to which value (f(x) or E(f(x)) to put at the top or bottom)
    plt.axvline(base_values, 0, ex_loc, color="#bbbbbb", linestyle="--", linewidth=0.5, zorder=-1)
    fx = base_values + values.sum()
    plt.axvline(fx, 0, fx_loc, color="#bbbbbb", linestyle="--", linewidth=0.5, zorder=-1)

    # clean up the main axis
    plt.gca().xaxis.set_ticks_position('bottom')
    plt.gca().yaxis.set_ticks_position('none')
    plt.gca().spines['right'].set_visible(False)
    plt.gca().spines['top'].set_visible(False)
    plt.gca().spines['left'].set_visible(False)
    ax.tick_params(labelsize=13)
    #plt.xlabel("\nModel output", fontsize=12)

    # Set up the two dictionaries depending on if f(x) or E(f(x)) is at top or bottom.
    # ax2_dict contains the details for the top, ax3_dict contains the details for the bottom.
    # The code used to place these two labels uses ax.twiny() which takes the opposite position of the given axis.
    # In order to switch the location of f(x) and E(f(x)) (by setting y_reverse=True) the code swaps the order 
    # that their respective axis is created.
    if y_reverse:
        # bottom of y axis: f(x)
        ax2_dict = {"value": base_values + values.sum(),
                    "ticklabels": ["$f(x)$", "$ = "+format_value(fx, "%0.03f")+"$"],        
                    "loc1": -10/72., "loc2": 12/72., "loc3": -15/72., "loc4": -15/72.}
        # top of y axis: E(f(x))
        ax3_dict = {"value": base_values, 
                    "ticklabels": ["\n$E[f(X)]$", "\n$ = "+format_value(base_values, "%0.03f")+"$"], 
                    "loc1": -20/72., "loc2": 22/72., "loc3": 0, "loc4": 0}
        
    else:
        # bottom of y axis: E(f(x))
        ax2_dict = {"value": base_values, 
                    "ticklabels": ["\n$E[f(X)]$", "\n$ = "+format_value(base_values, "%0.03f")+"$"], 
                    "loc1": -20/72., "loc2": 22/72., "loc3": 0, "loc4": -1/72.}
        # top of y axis: f(x)
        ax3_dict = {"value": base_values + values.sum(), 
                    "ticklabels": ["$f(x)$", "$ = "+format_value(fx, "%0.03f")+"$"], 
                    "loc1": -10/72., "loc2": 12/72., "loc3": 0, "loc4": 0}

    # draw the bottom tick mark (either f(X) or E(f(x)) depending on y_reverse)
    xmin, xmax = ax.get_xlim()
    ax2 = ax.twiny()
    ax2.set_xlim(xmin, xmax)
    # The 1e-8 is so matplotlib 3.3 doesn't try and collapse the ticks
    # provide two locations as using two tick labels (one for "f(x)" and one for "= value")
    ax2.set_xticks([ax2_dict["value"], ax2_dict["value"] + 1e-8])
    ax2.set_xticklabels(ax2_dict["ticklabels"], fontsize=12, ha="left")
    ax2.spines['right'].set_visible(False)
    ax2.spines['top'].set_visible(False)
    ax2.spines['left'].set_visible(False)

    # draw the top tick mark (either f(X) or E(f(x)) depending on y_reverse)
    ax3 = ax2.twiny()
    ax3.set_xlim(xmin, xmax)
    # The 1e-8 is so matplotlib 3.3 doesn't try and collapse the ticks
    # provide two locations as using two tick labels (one for "f(x)" and one for "= value"
    ax3.set_xticks([ax3_dict["value"], ax3_dict["value"] + 1e-8])  
    ax3.set_xticklabels(ax3_dict["ticklabels"], fontsize=12, ha="left")
    # adjust the position of the bottom label (either f(X) or E(f(x)) depending on y_reverse)
    tick_labels = ax3.xaxis.get_majorticklabels()
    tick_labels[0].set_transform(tick_labels[0].get_transform(
    ) + matplotlib.transforms.ScaledTranslation(ax3_dict["loc1"], ax3_dict["loc3"], fig.dpi_scale_trans))
    tick_labels[1].set_transform(tick_labels[1].get_transform(
    ) + matplotlib.transforms.ScaledTranslation(ax3_dict["loc2"], ax3_dict["loc4"], fig.dpi_scale_trans))
    tick_labels[1].set_color("#999999")
    ax3.spines['right'].set_visible(False)
    ax3.spines['top'].set_visible(False)
    ax3.spines['left'].set_visible(False)

    # adjust the position of the top label (either f(X) or E(f(x)) depending on y_reverse)
    tick_labels = ax2.xaxis.get_majorticklabels()
    tick_labels[0].set_transform(tick_labels[0].get_transform(
    ) + matplotlib.transforms.ScaledTranslation(ax2_dict["loc1"], ax2_dict["loc3"], fig.dpi_scale_trans))
    tick_labels[1].set_transform(tick_labels[1].get_transform(
    ) + matplotlib.transforms.ScaledTranslation(ax2_dict["loc2"], ax2_dict["loc4"], fig.dpi_scale_trans))
    tick_labels[1].set_color("#999999")

    # color the y tick labels that have the feature values as gray
    # (these fall behind the black ones with just the feature name)
    tick_labels = ax.yaxis.get_majorticklabels()
    for i in range(num_features):
        tick_labels[i].set_color("#999999")
            
    if y_reverse:
        # Reverse y axis (expected at top, predicted at bottom)
        ylim = ax.get_ylim() # Get current min and max
        ax.set_ylim(ylim[1], ylim[0]) # Reverse min and max
        
    if show:
        plt.show()
    else:
        return plt.gcf()
7f0c9b0fcaac689cc48c9067669e9f5b7708eaba
<mask token>

<mask token>


def waterfall(shap_values, max_display=10, show=True, y_reverse=False,
    rank_absolute=True, raw_ascending=True):
    """ Plots an explantion of a single prediction as a waterfall plot.
    The SHAP value of a feature represents the impact of the evidence provided by that feature on the model's
    output. The waterfall plot is designed to visually display how the SHAP values (evidence) of each feature
    move the model output from our prior expectation under the background data distribution, to the final model
    prediction given the evidence of all the features. Features are sorted by the magnitude of their SHAP values
    with the smallest magnitude features grouped together at the bottom of the plot when the number of features
    in the models exceeds the max_display parameter.
    Parameters
    ----------
    shap_values : Explanation
        A one-dimensional Explanation object that contains the feature values and SHAP values to plot.
    max_display : str
        The maximum number of features to plot.
    show : bool
        Whether matplotlib.pyplot.show() is called before returning. Setting this to False allows the plot
        to be customized further after it has been created.
    y_reverse : bool
        Determines placement of expected value & predicted value on y axis. Setting this to True reverses the 
        y axis and puts expected value at top of y axis, and predicted value at bottom of y axis.
    rank_absolute : bool
        If True then rank values by their absolute values.
    raw_ascending : bool
        If rank_absolute == False then look to this setting, else redundant (as order is on absolute values).
        Determines order of ranked features if not absolute. If True then from expected value put features in 
        order of most negative to least negative to smallest positive to most positive. If False then from expected 
        value put features in order of most positive to smallest positive to least negative to most negative.
    """
    if show is False:
        plt.ioff()
    base_values = shap_values.base_values
    features = (shap_values.display_data if shap_values.display_data is not
        None else shap_values.data)
    feature_names = shap_values.feature_names
    lower_bounds = getattr(shap_values, 'lower_bounds', None)
    upper_bounds = getattr(shap_values, 'upper_bounds', None)
    values = shap_values.values
    if type(base_values) == np.ndarray and len(base_values) > 0 or type(
        base_values) == list:
        raise Exception(
            'waterfall_plot requires a scalar base_values of the model output as the first parameter, but you have passed an array as the first parameter! Try shap.waterfall_plot(explainer.base_values[0], values[0], X[0]) or for multi-output models try shap.waterfall_plot(explainer.base_values[0], values[0][0], X[0]).'
            )
    if len(values.shape) == 2:
        raise Exception(
            'The waterfall_plot can currently only plot a single explanation but a matrix of explanations was passed!'
            )
    if safe_isinstance(features, 'pandas.core.series.Series'):
        if feature_names is None:
            feature_names = list(features.index)
        features = features.values
    if feature_names is None:
        feature_names = np.array([(labels['FEATURE'] % str(i)) for i in
            range(len(values))])
    num_features = min(max_display, len(values))
    row_height = 0.5
    rng = range(num_features - 1, -1, -1)
    order = np.argsort(-np.abs(values))
    pos_lefts = []
    pos_inds = []
    pos_widths = []
    pos_low = []
    pos_high = []
    neg_lefts = []
    neg_inds = []
    neg_widths = []
    neg_low = []
    neg_high = []
    loc = base_values + values.sum()
    yticklabels = ['' for i in range(num_features + 1)]
    plt.gcf().set_size_inches(8, num_features * row_height + 1.5)
    if num_features == len(values):
        num_individual = num_features
    else:
        num_individual = num_features - 1
    if not rank_absolute:
        if raw_ascending:
            values_to_sort = [values[order[i]] for i in range(num_individual)]
            arg_to_sort = order[0:num_individual]
            arg_values_to_sort = np.argsort(values_to_sort)
            order_arg_to_sort = [arg_to_sort[arg_values_to_sort[i]] for i in
                range(num_individual - 1, -1, -1)]
            order[0:num_individual] = order_arg_to_sort[0:num_individual]
        else:
            values_to_sort = [(-values[order[i]]) for i in range(
                num_individual)]
            arg_to_sort = order[0:num_individual]
            arg_values_to_sort = np.argsort(values_to_sort)
            order_arg_to_sort = [arg_to_sort[arg_values_to_sort[i]] for i in
                range(num_individual - 1, -1, -1)]
            order[0:num_individual] = order_arg_to_sort[0:num_individual]
    for i in range(num_individual):
        sval = values[order[i]]
        loc -= sval
        if sval >= 0:
            pos_inds.append(rng[i])
            pos_widths.append(sval)
            if lower_bounds is not None:
                pos_low.append(lower_bounds[order[i]])
                pos_high.append(upper_bounds[order[i]])
            pos_lefts.append(loc)
        else:
            neg_inds.append(rng[i])
            neg_widths.append(sval)
            if lower_bounds is not None:
                neg_low.append(lower_bounds[order[i]])
                neg_high.append(upper_bounds[order[i]])
            neg_lefts.append(loc)
        if num_individual != num_features or i + 4 < num_individual:
            plt.plot([loc, loc], [rng[i] - 1 - 0.4, rng[i] + 0.4], color=
                '#bbbbbb', linestyle='--', linewidth=0.5, zorder=-1)
        if features is None:
            yticklabels[rng[i]] = feature_names[order[i]]
        elif np.issubdtype(type(features[order[i]]), np.number):
            yticklabels[rng[i]] = format_value(float(features[order[i]]),
                '%0.03f') + ' = ' + feature_names[order[i]]
        else:
            yticklabels[rng[i]] = features[order[i]] + ' = ' + feature_names[
                order[i]]
    if num_features < len(values):
        yticklabels[0] = '%d other features' % (len(values) - num_features + 1)
        remaining_impact = base_values - loc
        if remaining_impact < 0:
            pos_inds.append(0)
            pos_widths.append(-remaining_impact)
            pos_lefts.append(loc + remaining_impact)
            c = colors.red_rgb
        else:
            neg_inds.append(0)
            neg_widths.append(-remaining_impact)
            neg_lefts.append(loc + remaining_impact)
            c = colors.blue_rgb
    points = pos_lefts + list(np.array(pos_lefts) + np.array(pos_widths)
        ) + neg_lefts + list(np.array(neg_lefts) + np.array(neg_widths))
    dataw = np.max(points) - np.min(points)
    label_padding = np.array([(0.1 * dataw if w < 1 else 0) for w in
        pos_widths])
    plt.barh(pos_inds, np.array(pos_widths) + label_padding + 0.02 * dataw,
        left=np.array(pos_lefts) - 0.01 * dataw, color=colors.red_rgb, alpha=0)
    label_padding = np.array([(-0.1 * dataw if -w < 1 else 0) for w in
        neg_widths])
    plt.barh(neg_inds, np.array(neg_widths) + label_padding - 0.02 * dataw,
        left=np.array(neg_lefts) + 0.01 * dataw, color=colors.blue_rgb, alpha=0
        )
    head_length = 0.08
    bar_width = 0.8
    xlen = plt.xlim()[1] - plt.xlim()[0]
    fig = plt.gcf()
    ax = plt.gca()
    xticks = ax.get_xticks()
    bbox = ax.get_window_extent().transformed(fig.dpi_scale_trans.inverted())
    width, height = bbox.width, bbox.height
    bbox_to_xscale = xlen / width
    hl_scaled = bbox_to_xscale * head_length
    renderer = fig.canvas.get_renderer()
    for i in range(len(pos_inds)):
        dist = pos_widths[i]
        arrow_obj = plt.arrow(pos_lefts[i], pos_inds[i], max(dist -
            hl_scaled, 1e-06), 0, head_length=min(dist, hl_scaled), color=
            colors.red_rgb, width=bar_width, head_width=bar_width)
        if pos_low is not None and i < len(pos_low):
            plt.errorbar(pos_lefts[i] + pos_widths[i], pos_inds[i], xerr=np
                .array([[pos_widths[i] - pos_low[i]], [pos_high[i] -
                pos_widths[i]]]), ecolor=colors.light_red_rgb)
        txt_obj = plt.text(pos_lefts[i] + 0.5 * dist, pos_inds[i],
            format_value(pos_widths[i], '%+0.02f'), horizontalalignment=
            'center', verticalalignment='center', color='white', fontsize=12)
        text_bbox = txt_obj.get_window_extent(renderer=renderer)
        arrow_bbox = arrow_obj.get_window_extent(renderer=renderer)
        if text_bbox.width > arrow_bbox.width:
            txt_obj.remove()
            txt_obj = plt.text(pos_lefts[i] + 5 / 72 * bbox_to_xscale +
                dist, pos_inds[i], format_value(pos_widths[i], '%+0.02f'),
                horizontalalignment='left', verticalalignment='center',
                color=colors.red_rgb, fontsize=12)
    for i in range(len(neg_inds)):
        dist = neg_widths[i]
        arrow_obj = plt.arrow(neg_lefts[i], neg_inds[i], -max(-dist -
            hl_scaled, 1e-06), 0, head_length=min(-dist, hl_scaled), color=
            colors.blue_rgb, width=bar_width, head_width=bar_width)
        if neg_low is not None and i < len(neg_low):
            plt.errorbar(neg_lefts[i] + neg_widths[i], neg_inds[i], xerr=np
                .array([[neg_widths[i] - neg_low[i]], [neg_high[i] -
                neg_widths[i]]]), ecolor=colors.light_blue_rgb)
        txt_obj = plt.text(neg_lefts[i] + 0.5 * dist, neg_inds[i],
            format_value(neg_widths[i], '%+0.02f'), horizontalalignment=
            'center', verticalalignment='center', color='white', fontsize=12)
        text_bbox = txt_obj.get_window_extent(renderer=renderer)
        arrow_bbox = arrow_obj.get_window_extent(renderer=renderer)
        if text_bbox.width > arrow_bbox.width:
            txt_obj.remove()
            txt_obj = plt.text(neg_lefts[i] - 5 / 72 * bbox_to_xscale +
                dist, neg_inds[i], format_value(neg_widths[i], '%+0.02f'),
                horizontalalignment='right', verticalalignment='center',
                color=colors.blue_rgb, fontsize=12)
    ytick_pos = list(range(num_features)) + list(np.arange(num_features) + 
        1e-08)
    plt.yticks(ytick_pos, yticklabels[:-1] + [l.split('=')[-1] for l in
        yticklabels[:-1]], fontsize=13)
    for i in range(num_features):
        plt.axhline(i, color='#cccccc', lw=0.5, dashes=(1, 5), zorder=-1)
    if y_reverse:
        ex_loc = 1
        fx_loc = 1 / num_features
    else:
        ex_loc = 1 / num_features
        fx_loc = 1
    plt.axvline(base_values, 0, ex_loc, color='#bbbbbb', linestyle='--',
        linewidth=0.5, zorder=-1)
    fx = base_values + values.sum()
    plt.axvline(fx, 0, fx_loc, color='#bbbbbb', linestyle='--', linewidth=
        0.5, zorder=-1)
    plt.gca().xaxis.set_ticks_position('bottom')
    plt.gca().yaxis.set_ticks_position('none')
    plt.gca().spines['right'].set_visible(False)
    plt.gca().spines['top'].set_visible(False)
    plt.gca().spines['left'].set_visible(False)
    ax.tick_params(labelsize=13)
    if y_reverse:
        ax2_dict = {'value': base_values + values.sum(), 'ticklabels': [
            '$f(x)$', '$ = ' + format_value(fx, '%0.03f') + '$'], 'loc1': -
            10 / 72.0, 'loc2': 12 / 72.0, 'loc3': -15 / 72.0, 'loc4': -15 /
            72.0}
        ax3_dict = {'value': base_values, 'ticklabels': ['\n$E[f(X)]$', 
            '\n$ = ' + format_value(base_values, '%0.03f') + '$'], 'loc1': 
            -20 / 72.0, 'loc2': 22 / 72.0, 'loc3': 0, 'loc4': 0}
    else:
        ax2_dict = {'value': base_values, 'ticklabels': ['\n$E[f(X)]$', 
            '\n$ = ' + format_value(base_values, '%0.03f') + '$'], 'loc1': 
            -20 / 72.0, 'loc2': 22 / 72.0, 'loc3': 0, 'loc4': -1 / 72.0}
        ax3_dict = {'value': base_values + values.sum(), 'ticklabels': [
            '$f(x)$', '$ = ' + format_value(fx, '%0.03f') + '$'], 'loc1': -
            10 / 72.0, 'loc2': 12 / 72.0, 'loc3': 0, 'loc4': 0}
    xmin, xmax = ax.get_xlim()
    ax2 = ax.twiny()
    ax2.set_xlim(xmin, xmax)
    ax2.set_xticks([ax2_dict['value'], ax2_dict['value'] + 1e-08])
    ax2.set_xticklabels(ax2_dict['ticklabels'], fontsize=12, ha='left')
    ax2.spines['right'].set_visible(False)
    ax2.spines['top'].set_visible(False)
    ax2.spines['left'].set_visible(False)
    ax3 = ax2.twiny()
    ax3.set_xlim(xmin, xmax)
    ax3.set_xticks([ax3_dict['value'], ax3_dict['value'] + 1e-08])
    ax3.set_xticklabels(ax3_dict['ticklabels'], fontsize=12, ha='left')
    tick_labels = ax3.xaxis.get_majorticklabels()
    tick_labels[0].set_transform(tick_labels[0].get_transform() +
        matplotlib.transforms.ScaledTranslation(ax3_dict['loc1'], ax3_dict[
        'loc3'], fig.dpi_scale_trans))
    tick_labels[1].set_transform(tick_labels[1].get_transform() +
        matplotlib.transforms.ScaledTranslation(ax3_dict['loc2'], ax3_dict[
        'loc4'], fig.dpi_scale_trans))
    tick_labels[1].set_color('#999999')
    ax3.spines['right'].set_visible(False)
    ax3.spines['top'].set_visible(False)
    ax3.spines['left'].set_visible(False)
    tick_labels = ax2.xaxis.get_majorticklabels()
    tick_labels[0].set_transform(tick_labels[0].get_transform() +
        matplotlib.transforms.ScaledTranslation(ax2_dict['loc1'], ax2_dict[
        'loc3'], fig.dpi_scale_trans))
    tick_labels[1].set_transform(tick_labels[1].get_transform() +
        matplotlib.transforms.ScaledTranslation(ax2_dict['loc2'], ax2_dict[
        'loc4'], fig.dpi_scale_trans))
    tick_labels[1].set_color('#999999')
    tick_labels = ax.yaxis.get_majorticklabels()
    for i in range(num_features):
        tick_labels[i].set_color('#999999')
    if y_reverse:
        ylim = ax.get_ylim()
        ax.set_ylim(ylim[1], ylim[0])
    if show:
        plt.show()
    else:
        return plt.gcf()

<mask token>
try:
    import matplotlib.pyplot as plt
    import matplotlib
except ImportError:
    warnings.warn('matplotlib could not be loaded!')
    pass
<mask token>


def waterfall(shap_values, max_display=10, show=True, y_reverse=False,
    rank_absolute=True, raw_ascending=True):
    """ Plots an explantion of a single prediction as a waterfall plot.
    The SHAP value of a feature represents the impact of the evidence provided by that feature on the model's
    output. The waterfall plot is designed to visually display how the SHAP values (evidence) of each feature
    move the model output from our prior expectation under the background data distribution, to the final model
    prediction given the evidence of all the features. Features are sorted by the magnitude of their SHAP values
    with the smallest magnitude features grouped together at the bottom of the plot when the number of features
    in the models exceeds the max_display parameter.
    Parameters
    ----------
    shap_values : Explanation
        A one-dimensional Explanation object that contains the feature values and SHAP values to plot.
    max_display : str
        The maximum number of features to plot.
    show : bool
        Whether matplotlib.pyplot.show() is called before returning. Setting this to False allows the plot
        to be customized further after it has been created.
    y_reverse : bool
        Determines placement of expected value & predicted value on y axis. Setting this to True reverses the 
        y axis and puts expected value at top of y axis, and predicted value at bottom of y axis.
    rank_absolute : bool
        If True then rank values by their absolute values.
    raw_ascending : bool
        If rank_absolute == False then look to this setting, else redundant (as order is on absolute values).
        Determines order of ranked features if not absolute. If True then from expected value put features in 
        order of most negative to least negative to smallest positive to most positive. If False then from expected 
        value put features in order of most positive to smallest positive to least negative to most negative.
    """
    if show is False:
        plt.ioff()
    base_values = shap_values.base_values
    features = (shap_values.display_data if shap_values.display_data is not
        None else shap_values.data)
    feature_names = shap_values.feature_names
    lower_bounds = getattr(shap_values, 'lower_bounds', None)
    upper_bounds = getattr(shap_values, 'upper_bounds', None)
    values = shap_values.values
    if type(base_values) == np.ndarray and len(base_values) > 0 or type(
        base_values) == list:
        raise Exception(
            'waterfall_plot requires a scalar base_values of the model output as the first parameter, but you have passed an array as the first parameter! Try shap.waterfall_plot(explainer.base_values[0], values[0], X[0]) or for multi-output models try shap.waterfall_plot(explainer.base_values[0], values[0][0], X[0]).'
            )
    if len(values.shape) == 2:
        raise Exception(
            'The waterfall_plot can currently only plot a single explanation but a matrix of explanations was passed!'
            )
    if safe_isinstance(features, 'pandas.core.series.Series'):
        if feature_names is None:
            feature_names = list(features.index)
        features = features.values
    if feature_names is None:
        feature_names = np.array([(labels['FEATURE'] % str(i)) for i in
            range(len(values))])
    num_features = min(max_display, len(values))
    row_height = 0.5
    rng = range(num_features - 1, -1, -1)
    order = np.argsort(-np.abs(values))
    pos_lefts = []
    pos_inds = []
    pos_widths = []
    pos_low = []
    pos_high = []
    neg_lefts = []
    neg_inds = []
    neg_widths = []
    neg_low = []
    neg_high = []
    loc = base_values + values.sum()
    yticklabels = ['' for i in range(num_features + 1)]
    plt.gcf().set_size_inches(8, num_features * row_height + 1.5)
    if num_features == len(values):
        num_individual = num_features
    else:
        num_individual = num_features - 1
    if not rank_absolute:
        if raw_ascending:
            values_to_sort = [values[order[i]] for i in range(num_individual)]
            arg_to_sort = order[0:num_individual]
            arg_values_to_sort = np.argsort(values_to_sort)
            order_arg_to_sort = [arg_to_sort[arg_values_to_sort[i]] for i in
                range(num_individual - 1, -1, -1)]
            order[0:num_individual] = order_arg_to_sort[0:num_individual]
        else:
            values_to_sort = [(-values[order[i]]) for i in range(
                num_individual)]
            arg_to_sort = order[0:num_individual]
            arg_values_to_sort = np.argsort(values_to_sort)
            order_arg_to_sort = [arg_to_sort[arg_values_to_sort[i]] for i in
                range(num_individual - 1, -1, -1)]
            order[0:num_individual] = order_arg_to_sort[0:num_individual]
    for i in range(num_individual):
        sval = values[order[i]]
        loc -= sval
        if sval >= 0:
            pos_inds.append(rng[i])
            pos_widths.append(sval)
            if lower_bounds is not None:
                pos_low.append(lower_bounds[order[i]])
                pos_high.append(upper_bounds[order[i]])
            pos_lefts.append(loc)
        else:
            neg_inds.append(rng[i])
            neg_widths.append(sval)
            if lower_bounds is not None:
                neg_low.append(lower_bounds[order[i]])
                neg_high.append(upper_bounds[order[i]])
            neg_lefts.append(loc)
        if num_individual != num_features or i + 4 < num_individual:
            plt.plot([loc, loc], [rng[i] - 1 - 0.4, rng[i] + 0.4], color=
                '#bbbbbb', linestyle='--', linewidth=0.5, zorder=-1)
        if features is None:
            yticklabels[rng[i]] = feature_names[order[i]]
        elif np.issubdtype(type(features[order[i]]), np.number):
            yticklabels[rng[i]] = format_value(float(features[order[i]]),
                '%0.03f') + ' = ' + feature_names[order[i]]
        else:
            yticklabels[rng[i]] = features[order[i]] + ' = ' + feature_names[
                order[i]]
    if num_features < len(values):
        yticklabels[0] = '%d other features' % (len(values) - num_features + 1)
        remaining_impact = base_values - loc
        if remaining_impact < 0:
            pos_inds.append(0)
            pos_widths.append(-remaining_impact)
            pos_lefts.append(loc + remaining_impact)
            c = colors.red_rgb
        else:
            neg_inds.append(0)
            neg_widths.append(-remaining_impact)
            neg_lefts.append(loc + remaining_impact)
            c = colors.blue_rgb
    points = pos_lefts + list(np.array(pos_lefts) + np.array(pos_widths)
        ) + neg_lefts + list(np.array(neg_lefts) + np.array(neg_widths))
    dataw = np.max(points) - np.min(points)
    label_padding = np.array([(0.1 * dataw if w < 1 else 0) for w in
        pos_widths])
    plt.barh(pos_inds, np.array(pos_widths) + label_padding + 0.02 * dataw,
        left=np.array(pos_lefts) - 0.01 * dataw, color=colors.red_rgb, alpha=0)
    label_padding = np.array([(-0.1 * dataw if -w < 1 else 0) for w in
        neg_widths])
    plt.barh(neg_inds, np.array(neg_widths) + label_padding - 0.02 * dataw,
        left=np.array(neg_lefts) + 0.01 * dataw, color=colors.blue_rgb, alpha=0
        )
    head_length = 0.08
    bar_width = 0.8
    xlen = plt.xlim()[1] - plt.xlim()[0]
    fig = plt.gcf()
    ax = plt.gca()
    xticks = ax.get_xticks()
    bbox = ax.get_window_extent().transformed(fig.dpi_scale_trans.inverted())
    width, height = bbox.width, bbox.height
    bbox_to_xscale = xlen / width
    hl_scaled = bbox_to_xscale * head_length
    renderer = fig.canvas.get_renderer()
    for i in range(len(pos_inds)):
        dist = pos_widths[i]
        arrow_obj = plt.arrow(pos_lefts[i], pos_inds[i], max(dist -
            hl_scaled, 1e-06), 0, head_length=min(dist, hl_scaled), color=
            colors.red_rgb, width=bar_width, head_width=bar_width)
        if pos_low is not None and i < len(pos_low):
            plt.errorbar(pos_lefts[i] + pos_widths[i], pos_inds[i], xerr=np
                .array([[pos_widths[i] - pos_low[i]], [pos_high[i] -
                pos_widths[i]]]), ecolor=colors.light_red_rgb)
        txt_obj = plt.text(pos_lefts[i] + 0.5 * dist, pos_inds[i],
            format_value(pos_widths[i], '%+0.02f'), horizontalalignment=
            'center', verticalalignment='center', color='white', fontsize=12)
        text_bbox = txt_obj.get_window_extent(renderer=renderer)
        arrow_bbox = arrow_obj.get_window_extent(renderer=renderer)
        if text_bbox.width > arrow_bbox.width:
            txt_obj.remove()
            txt_obj = plt.text(pos_lefts[i] + 5 / 72 * bbox_to_xscale +
                dist, pos_inds[i], format_value(pos_widths[i], '%+0.02f'),
                horizontalalignment='left', verticalalignment='center',
                color=colors.red_rgb, fontsize=12)
    for i in range(len(neg_inds)):
        dist = neg_widths[i]
        arrow_obj = plt.arrow(neg_lefts[i], neg_inds[i], -max(-dist -
            hl_scaled, 1e-06), 0, head_length=min(-dist, hl_scaled), color=
            colors.blue_rgb, width=bar_width, head_width=bar_width)
        if neg_low is not None and i < len(neg_low):
            plt.errorbar(neg_lefts[i] + neg_widths[i], neg_inds[i], xerr=np
                .array([[neg_widths[i] - neg_low[i]], [neg_high[i] -
                neg_widths[i]]]), ecolor=colors.light_blue_rgb)
        txt_obj = plt.text(neg_lefts[i] + 0.5 * dist, neg_inds[i],
            format_value(neg_widths[i], '%+0.02f'), horizontalalignment=
            'center', verticalalignment='center', color='white', fontsize=12)
        text_bbox = txt_obj.get_window_extent(renderer=renderer)
        arrow_bbox = arrow_obj.get_window_extent(renderer=renderer)
        if text_bbox.width > arrow_bbox.width:
            txt_obj.remove()
            txt_obj = plt.text(neg_lefts[i] - 5 / 72 * bbox_to_xscale +
                dist, neg_inds[i], format_value(neg_widths[i], '%+0.02f'),
                horizontalalignment='right', verticalalignment='center',
                color=colors.blue_rgb, fontsize=12)
    ytick_pos = list(range(num_features)) + list(np.arange(num_features) + 
        1e-08)
    plt.yticks(ytick_pos, yticklabels[:-1] + [l.split('=')[-1] for l in
        yticklabels[:-1]], fontsize=13)
    for i in range(num_features):
        plt.axhline(i, color='#cccccc', lw=0.5, dashes=(1, 5), zorder=-1)
    if y_reverse:
        ex_loc = 1
        fx_loc = 1 / num_features
    else:
        ex_loc = 1 / num_features
        fx_loc = 1
    plt.axvline(base_values, 0, ex_loc, color='#bbbbbb', linestyle='--',
        linewidth=0.5, zorder=-1)
    fx = base_values + values.sum()
    plt.axvline(fx, 0, fx_loc, color='#bbbbbb', linestyle='--', linewidth=
        0.5, zorder=-1)
    plt.gca().xaxis.set_ticks_position('bottom')
    plt.gca().yaxis.set_ticks_position('none')
    plt.gca().spines['right'].set_visible(False)
    plt.gca().spines['top'].set_visible(False)
    plt.gca().spines['left'].set_visible(False)
    ax.tick_params(labelsize=13)
    if y_reverse:
        ax2_dict = {'value': base_values + values.sum(), 'ticklabels': [
            '$f(x)$', '$ = ' + format_value(fx, '%0.03f') + '$'], 'loc1': -
            10 / 72.0, 'loc2': 12 / 72.0, 'loc3': -15 / 72.0, 'loc4': -15 /
            72.0}
        ax3_dict = {'value': base_values, 'ticklabels': ['\n$E[f(X)]$', 
            '\n$ = ' + format_value(base_values, '%0.03f') + '$'], 'loc1': 
            -20 / 72.0, 'loc2': 22 / 72.0, 'loc3': 0, 'loc4': 0}
    else:
        ax2_dict = {'value': base_values, 'ticklabels': ['\n$E[f(X)]$', 
            '\n$ = ' + format_value(base_values, '%0.03f') + '$'], 'loc1': 
            -20 / 72.0, 'loc2': 22 / 72.0, 'loc3': 0, 'loc4': -1 / 72.0}
        ax3_dict = {'value': base_values + values.sum(), 'ticklabels': [
            '$f(x)$', '$ = ' + format_value(fx, '%0.03f') + '$'], 'loc1': -
            10 / 72.0, 'loc2': 12 / 72.0, 'loc3': 0, 'loc4': 0}
    xmin, xmax = ax.get_xlim()
    ax2 = ax.twiny()
    ax2.set_xlim(xmin, xmax)
    ax2.set_xticks([ax2_dict['value'], ax2_dict['value'] + 1e-08])
    ax2.set_xticklabels(ax2_dict['ticklabels'], fontsize=12, ha='left')
    ax2.spines['right'].set_visible(False)
    ax2.spines['top'].set_visible(False)
    ax2.spines['left'].set_visible(False)
    ax3 = ax2.twiny()
    ax3.set_xlim(xmin, xmax)
    ax3.set_xticks([ax3_dict['value'], ax3_dict['value'] + 1e-08])
    ax3.set_xticklabels(ax3_dict['ticklabels'], fontsize=12, ha='left')
    tick_labels = ax3.xaxis.get_majorticklabels()
    tick_labels[0].set_transform(tick_labels[0].get_transform() +
        matplotlib.transforms.ScaledTranslation(ax3_dict['loc1'], ax3_dict[
        'loc3'], fig.dpi_scale_trans))
    tick_labels[1].set_transform(tick_labels[1].get_transform() +
        matplotlib.transforms.ScaledTranslation(ax3_dict['loc2'], ax3_dict[
        'loc4'], fig.dpi_scale_trans))
    tick_labels[1].set_color('#999999')
    ax3.spines['right'].set_visible(False)
    ax3.spines['top'].set_visible(False)
    ax3.spines['left'].set_visible(False)
    tick_labels = ax2.xaxis.get_majorticklabels()
    tick_labels[0].set_transform(tick_labels[0].get_transform() +
        matplotlib.transforms.ScaledTranslation(ax2_dict['loc1'], ax2_dict[
        'loc3'], fig.dpi_scale_trans))
    tick_labels[1].set_transform(tick_labels[1].get_transform() +
        matplotlib.transforms.ScaledTranslation(ax2_dict['loc2'], ax2_dict[
        'loc4'], fig.dpi_scale_trans))
    tick_labels[1].set_color('#999999')
    tick_labels = ax.yaxis.get_majorticklabels()
    for i in range(num_features):
        tick_labels[i].set_color('#999999')
    if y_reverse:
        ylim = ax.get_ylim()
        ax.set_ylim(ylim[1], ylim[0])
    if show:
        plt.show()
    else:
        return plt.gcf()

import numpy as np
import warnings
try:
    import matplotlib.pyplot as plt
    import matplotlib
except ImportError:
    warnings.warn('matplotlib could not be loaded!')
    pass
from .labels import labels
from .general import safe_isinstance, format_value
from . import colors


def waterfall(shap_values, max_display=10, show=True, y_reverse=False,
    rank_absolute=True, raw_ascending=True):
    """ Plots an explantion of a single prediction as a waterfall plot.
    The SHAP value of a feature represents the impact of the evidence provided by that feature on the model's
    output. The waterfall plot is designed to visually display how the SHAP values (evidence) of each feature
    move the model output from our prior expectation under the background data distribution, to the final model
    prediction given the evidence of all the features. Features are sorted by the magnitude of their SHAP values
    with the smallest magnitude features grouped together at the bottom of the plot when the number of features
    in the models exceeds the max_display parameter.
    Parameters
    ----------
    shap_values : Explanation
        A one-dimensional Explanation object that contains the feature values and SHAP values to plot.
    max_display : str
        The maximum number of features to plot.
    show : bool
        Whether matplotlib.pyplot.show() is called before returning. Setting this to False allows the plot
        to be customized further after it has been created.
    y_reverse : bool
        Determines placement of expected value & predicted value on y axis. Setting this to True reverses the 
        y axis and puts expected value at top of y axis, and predicted value at bottom of y axis.
    rank_absolute : bool
        If True then rank values by their absolute values.
    raw_ascending : bool
        If rank_absolute == False then look to this setting, else redundant (as order is on absolute values).
        Determines order of ranked features if not absolute. If True then from expected value put features in 
        order of most negative to least negative to smallest positive to most positive. If False then from expected 
        value put features in order of most positive to smallest positive to least negative to most negative.
    """
    if show is False:
        plt.ioff()
    base_values = shap_values.base_values
    features = (shap_values.display_data if shap_values.display_data is not
        None else shap_values.data)
    feature_names = shap_values.feature_names
    lower_bounds = getattr(shap_values, 'lower_bounds', None)
    upper_bounds = getattr(shap_values, 'upper_bounds', None)
    values = shap_values.values
    if type(base_values) == np.ndarray and len(base_values) > 0 or type(
        base_values) == list:
        raise Exception(
            'waterfall_plot requires a scalar base_values of the model output as the first parameter, but you have passed an array as the first parameter! Try shap.waterfall_plot(explainer.base_values[0], values[0], X[0]) or for multi-output models try shap.waterfall_plot(explainer.base_values[0], values[0][0], X[0]).'
            )
    if len(values.shape) == 2:
        raise Exception(
            'The waterfall_plot can currently only plot a single explanation but a matrix of explanations was passed!'
            )
    if safe_isinstance(features, 'pandas.core.series.Series'):
        if feature_names is None:
            feature_names = list(features.index)
        features = features.values
    if feature_names is None:
        feature_names = np.array([(labels['FEATURE'] % str(i)) for i in
            range(len(values))])
    num_features = min(max_display, len(values))
    row_height = 0.5
    rng = range(num_features - 1, -1, -1)
    order = np.argsort(-np.abs(values))
    pos_lefts = []
    pos_inds = []
    pos_widths = []
    pos_low = []
    pos_high = []
    neg_lefts = []
    neg_inds = []
    neg_widths = []
    neg_low = []
    neg_high = []
    loc = base_values + values.sum()
    yticklabels = ['' for i in range(num_features + 1)]
    plt.gcf().set_size_inches(8, num_features * row_height + 1.5)
    if num_features == len(values):
        num_individual = num_features
    else:
        num_individual = num_features - 1
    if not rank_absolute:
        if raw_ascending:
            values_to_sort = [values[order[i]] for i in range(num_individual)]
            arg_to_sort = order[0:num_individual]
            arg_values_to_sort = np.argsort(values_to_sort)
            order_arg_to_sort = [arg_to_sort[arg_values_to_sort[i]] for i in
                range(num_individual - 1, -1, -1)]
            order[0:num_individual] = order_arg_to_sort[0:num_individual]
        else:
            values_to_sort = [(-values[order[i]]) for i in range(
                num_individual)]
            arg_to_sort = order[0:num_individual]
            arg_values_to_sort = np.argsort(values_to_sort)
            order_arg_to_sort = [arg_to_sort[arg_values_to_sort[i]] for i in
                range(num_individual - 1, -1, -1)]
            order[0:num_individual] = order_arg_to_sort[0:num_individual]
    for i in range(num_individual):
        sval = values[order[i]]
        loc -= sval
        if sval >= 0:
            pos_inds.append(rng[i])
            pos_widths.append(sval)
            if lower_bounds is not None:
                pos_low.append(lower_bounds[order[i]])
                pos_high.append(upper_bounds[order[i]])
            pos_lefts.append(loc)
        else:
            neg_inds.append(rng[i])
            neg_widths.append(sval)
            if lower_bounds is not None:
                neg_low.append(lower_bounds[order[i]])
                neg_high.append(upper_bounds[order[i]])
            neg_lefts.append(loc)
        if num_individual != num_features or i + 4 < num_individual:
            plt.plot([loc, loc], [rng[i] - 1 - 0.4, rng[i] + 0.4], color=
                '#bbbbbb', linestyle='--', linewidth=0.5, zorder=-1)
        if features is None:
            yticklabels[rng[i]] = feature_names[order[i]]
        elif np.issubdtype(type(features[order[i]]), np.number):
            yticklabels[rng[i]] = format_value(float(features[order[i]]),
                '%0.03f') + ' = ' + feature_names[order[i]]
        else:
            yticklabels[rng[i]] = features[order[i]] + ' = ' + feature_names[
                order[i]]
    if num_features < len(values):
        yticklabels[0] = '%d other features' % (len(values) - num_features + 1)
        remaining_impact = base_values - loc
        if remaining_impact < 0:
            pos_inds.append(0)
            pos_widths.append(-remaining_impact)
            pos_lefts.append(loc + remaining_impact)
            c = colors.red_rgb
        else:
            neg_inds.append(0)
            neg_widths.append(-remaining_impact)
            neg_lefts.append(loc + remaining_impact)
            c = colors.blue_rgb
    points = pos_lefts + list(np.array(pos_lefts) + np.array(pos_widths)
        ) + neg_lefts + list(np.array(neg_lefts) + np.array(neg_widths))
    dataw = np.max(points) - np.min(points)
    label_padding = np.array([(0.1 * dataw if w < 1 else 0) for w in
        pos_widths])
    plt.barh(pos_inds, np.array(pos_widths) + label_padding + 0.02 * dataw,
        left=np.array(pos_lefts) - 0.01 * dataw, color=colors.red_rgb, alpha=0)
    label_padding = np.array([(-0.1 * dataw if -w < 1 else 0) for w in
        neg_widths])
    plt.barh(neg_inds, np.array(neg_widths) + label_padding - 0.02 * dataw,
        left=np.array(neg_lefts) + 0.01 * dataw, color=colors.blue_rgb, alpha=0
        )
    head_length = 0.08
    bar_width = 0.8
    xlen = plt.xlim()[1] - plt.xlim()[0]
    fig = plt.gcf()
    ax = plt.gca()
    xticks = ax.get_xticks()
    bbox = ax.get_window_extent().transformed(fig.dpi_scale_trans.inverted())
    width, height = bbox.width, bbox.height
    bbox_to_xscale = xlen / width
    hl_scaled = bbox_to_xscale * head_length
    renderer = fig.canvas.get_renderer()
    for i in range(len(pos_inds)):
        dist = pos_widths[i]
        arrow_obj = plt.arrow(pos_lefts[i], pos_inds[i], max(dist -
            hl_scaled, 1e-06), 0, head_length=min(dist, hl_scaled), color=
            colors.red_rgb, width=bar_width, head_width=bar_width)
        if pos_low is not None and i < len(pos_low):
            plt.errorbar(pos_lefts[i] + pos_widths[i], pos_inds[i], xerr=np
                .array([[pos_widths[i] - pos_low[i]], [pos_high[i] -
                pos_widths[i]]]), ecolor=colors.light_red_rgb)
        txt_obj = plt.text(pos_lefts[i] + 0.5 * dist, pos_inds[i],
            format_value(pos_widths[i], '%+0.02f'), horizontalalignment=
            'center', verticalalignment='center', color='white', fontsize=12)
        text_bbox = txt_obj.get_window_extent(renderer=renderer)
        arrow_bbox = arrow_obj.get_window_extent(renderer=renderer)
        if text_bbox.width > arrow_bbox.width:
            txt_obj.remove()
            txt_obj = plt.text(pos_lefts[i] + 5 / 72 * bbox_to_xscale +
                dist, pos_inds[i], format_value(pos_widths[i], '%+0.02f'),
                horizontalalignment='left', verticalalignment='center',
                color=colors.red_rgb, fontsize=12)
    for i in range(len(neg_inds)):
        dist = neg_widths[i]
        arrow_obj = plt.arrow(neg_lefts[i], neg_inds[i], -max(-dist -
            hl_scaled, 1e-06), 0, head_length=min(-dist, hl_scaled), color=
            colors.blue_rgb, width=bar_width, head_width=bar_width)
        if neg_low is not None and i < len(neg_low):
            plt.errorbar(neg_lefts[i] + neg_widths[i], neg_inds[i], xerr=np
                .array([[neg_widths[i] - neg_low[i]], [neg_high[i] -
                neg_widths[i]]]), ecolor=colors.light_blue_rgb)
        txt_obj = plt.text(neg_lefts[i] + 0.5 * dist, neg_inds[i],
            format_value(neg_widths[i], '%+0.02f'), horizontalalignment=
            'center', verticalalignment='center', color='white', fontsize=12)
        text_bbox = txt_obj.get_window_extent(renderer=renderer)
        arrow_bbox = arrow_obj.get_window_extent(renderer=renderer)
        if text_bbox.width > arrow_bbox.width:
            txt_obj.remove()
            txt_obj = plt.text(neg_lefts[i] - 5 / 72 * bbox_to_xscale +
                dist, neg_inds[i], format_value(neg_widths[i], '%+0.02f'),
                horizontalalignment='right', verticalalignment='center',
                color=colors.blue_rgb, fontsize=12)
    ytick_pos = list(range(num_features)) + list(np.arange(num_features) + 
        1e-08)
    plt.yticks(ytick_pos, yticklabels[:-1] + [l.split('=')[-1] for l in
        yticklabels[:-1]], fontsize=13)
    for i in range(num_features):
        plt.axhline(i, color='#cccccc', lw=0.5, dashes=(1, 5), zorder=-1)
    if y_reverse:
        ex_loc = 1
        fx_loc = 1 / num_features
    else:
        ex_loc = 1 / num_features
        fx_loc = 1
    plt.axvline(base_values, 0, ex_loc, color='#bbbbbb', linestyle='--',
        linewidth=0.5, zorder=-1)
    fx = base_values + values.sum()
    plt.axvline(fx, 0, fx_loc, color='#bbbbbb', linestyle='--', linewidth=
        0.5, zorder=-1)
    plt.gca().xaxis.set_ticks_position('bottom')
    plt.gca().yaxis.set_ticks_position('none')
    plt.gca().spines['right'].set_visible(False)
    plt.gca().spines['top'].set_visible(False)
    plt.gca().spines['left'].set_visible(False)
    ax.tick_params(labelsize=13)
    if y_reverse:
        ax2_dict = {'value': base_values + values.sum(), 'ticklabels': [
            '$f(x)$', '$ = ' + format_value(fx, '%0.03f') + '$'], 'loc1': -
            10 / 72.0, 'loc2': 12 / 72.0, 'loc3': -15 / 72.0, 'loc4': -15 /
            72.0}
        ax3_dict = {'value': base_values, 'ticklabels': ['\n$E[f(X)]$', 
            '\n$ = ' + format_value(base_values, '%0.03f') + '$'], 'loc1': 
            -20 / 72.0, 'loc2': 22 / 72.0, 'loc3': 0, 'loc4': 0}
    else:
        ax2_dict = {'value': base_values, 'ticklabels': ['\n$E[f(X)]$', 
            '\n$ = ' + format_value(base_values, '%0.03f') + '$'], 'loc1': 
            -20 / 72.0, 'loc2': 22 / 72.0, 'loc3': 0, 'loc4': -1 / 72.0}
        ax3_dict = {'value': base_values + values.sum(), 'ticklabels': [
            '$f(x)$', '$ = ' + format_value(fx, '%0.03f') + '$'], 'loc1': -
            10 / 72.0, 'loc2': 12 / 72.0, 'loc3': 0, 'loc4': 0}
    xmin, xmax = ax.get_xlim()
    ax2 = ax.twiny()
    ax2.set_xlim(xmin, xmax)
    ax2.set_xticks([ax2_dict['value'], ax2_dict['value'] + 1e-08])
    ax2.set_xticklabels(ax2_dict['ticklabels'], fontsize=12, ha='left')
    ax2.spines['right'].set_visible(False)
    ax2.spines['top'].set_visible(False)
    ax2.spines['left'].set_visible(False)
    ax3 = ax2.twiny()
    ax3.set_xlim(xmin, xmax)
    ax3.set_xticks([ax3_dict['value'], ax3_dict['value'] + 1e-08])
    ax3.set_xticklabels(ax3_dict['ticklabels'], fontsize=12, ha='left')
    tick_labels = ax3.xaxis.get_majorticklabels()
    tick_labels[0].set_transform(tick_labels[0].get_transform() +
        matplotlib.transforms.ScaledTranslation(ax3_dict['loc1'], ax3_dict[
        'loc3'], fig.dpi_scale_trans))
    tick_labels[1].set_transform(tick_labels[1].get_transform() +
        matplotlib.transforms.ScaledTranslation(ax3_dict['loc2'], ax3_dict[
        'loc4'], fig.dpi_scale_trans))
    tick_labels[1].set_color('#999999')
    ax3.spines['right'].set_visible(False)
    ax3.spines['top'].set_visible(False)
    ax3.spines['left'].set_visible(False)
    tick_labels = ax2.xaxis.get_majorticklabels()
    tick_labels[0].set_transform(tick_labels[0].get_transform() +
        matplotlib.transforms.ScaledTranslation(ax2_dict['loc1'], ax2_dict[
        'loc3'], fig.dpi_scale_trans))
    tick_labels[1].set_transform(tick_labels[1].get_transform() +
        matplotlib.transforms.ScaledTranslation(ax2_dict['loc2'], ax2_dict[
        'loc4'], fig.dpi_scale_trans))
    tick_labels[1].set_color('#999999')
    tick_labels = ax.yaxis.get_majorticklabels()
    for i in range(num_features):
        tick_labels[i].set_color('#999999')
    if y_reverse:
        ylim = ax.get_ylim()
        ax.set_ylim(ylim[1], ylim[0])
    if show:
        plt.show()
    else:
        return plt.gcf()

import numpy as np
import warnings
try:
    import matplotlib.pyplot as plt
    import matplotlib
except ImportError:
    warnings.warn("matplotlib could not be loaded!")
    pass
from .labels import labels
from .general import safe_isinstance, format_value
from . import colors

# 221116 Added argument (y_reverse) to function waterfall that reverses the y axis when set to True
# (y_reverse = True: Expected value at top of y axis, Predicted value at bottom of y axis)

# TODO: If we make a JS version of this plot then we could let users click on a bar and then see the dependence
# plot that is associated with that feature get overlayed on the plot...it would quickly allow users to answer
# why a feature is pushing down or up. Perhaps the best way to do this would be with an ICE plot hanging off
# of the bar...
def waterfall(shap_values, max_display=10, show=True, y_reverse=False, rank_absolute=True, raw_ascending=True):
    """ Plots an explantion of a single prediction as a waterfall plot.
    The SHAP value of a feature represents the impact of the evidence provided by that feature on the model's
    output. The waterfall plot is designed to visually display how the SHAP values (evidence) of each feature
    move the model output from our prior expectation under the background data distribution, to the final model
    prediction given the evidence of all the features. Features are sorted by the magnitude of their SHAP values
    with the smallest magnitude features grouped together at the bottom of the plot when the number of features
    in the models exceeds the max_display parameter.
    Parameters
    ----------
    shap_values : Explanation
        A one-dimensional Explanation object that contains the feature values and SHAP values to plot.
    max_display : str
        The maximum number of features to plot.
    show : bool
        Whether matplotlib.pyplot.show() is called before returning. Setting this to False allows the plot
        to be customized further after it has been created.
    y_reverse : bool
        Determines placement of expected value & predicted value on y axis. Setting this to True reverses the 
        y axis and puts expected value at top of y axis, and predicted value at bottom of y axis.
    rank_absolute : bool
        If True then rank values by their absolute values.
    raw_ascending : bool
        If rank_absolute == False then look to this setting, else redundant (as order is on absolute values).
        Determines order of ranked features if not absolute. If True then from expected value put features in 
        order of most negative to least negative to smallest positive to most positive. If False then from expected 
        value put features in order of most positive to smallest positive to least negative to most negative.
    """
    # Turn off interactive plot
    if show is False:
        plt.ioff()

    base_values = shap_values.base_values
    features = shap_values.display_data if shap_values.display_data is not None else shap_values.data
    feature_names = shap_values.feature_names
    lower_bounds = getattr(shap_values, "lower_bounds", None)
    upper_bounds = getattr(shap_values, "upper_bounds", None)
    values = shap_values.values

    # make sure we only have a single output to explain
    if (type(base_values) == np.ndarray and len(base_values) > 0) or type(base_values) == list:
        raise Exception("waterfall_plot requires a scalar base_values of the model output as the first "
                        "parameter, but you have passed an array as the first parameter! "
                        "Try shap.waterfall_plot(explainer.base_values[0], values[0], X[0]) or "
                        "for multi-output models try "
                        "shap.waterfall_plot(explainer.base_values[0], values[0][0], X[0]).")

    # make sure we only have a single explanation to plot
    if len(values.shape) == 2:
        raise Exception(
            "The waterfall_plot can currently only plot a single explanation but a matrix of explanations was passed!")

    # unwrap pandas series
    if safe_isinstance(features, "pandas.core.series.Series"):
        if feature_names is None:
            feature_names = list(features.index)
        features = features.values

    # fallback feature names
    if feature_names is None:
        feature_names = np.array([labels['FEATURE'] % str(i) for i in range(len(values))])

    # init variables we use for tracking the plot locations
    num_features = min(max_display, len(values))
    row_height = 0.5
    rng = range(num_features - 1, -1, -1)
    
    # Rank features by absolute values to choose which are going to be included individually in the 
    #   waterfall plot, the rest will be represented in one summed bar.
    order = np.argsort(-np.abs(values))

    pos_lefts = []
    pos_inds = []
    pos_widths = []
    pos_low = []
    pos_high = []
    neg_lefts = []
    neg_inds = []
    neg_widths = []
    neg_low = []
    neg_high = []
    loc = base_values + values.sum()
    yticklabels = ["" for i in range(num_features + 1)]

    # size the plot based on how many features we are plotting
    plt.gcf().set_size_inches(8, num_features * row_height + 1.5)

    # see how many individual (vs. grouped at the end) features we are plotting
    if num_features == len(values):
        num_individual = num_features
    else:
        num_individual = num_features - 1
        
    if not rank_absolute:
        # The plotting order of just the features that are going to be shown indiviually to use raw values
        if raw_ascending:
            # in ascending order (most negative to minimal change to most positive)
            values_to_sort = [values[order[i]] for i in range(num_individual)]
            arg_to_sort = order[0:num_individual]
            arg_values_to_sort = np.argsort(values_to_sort)
            # range(num_individual-1,-1,-1) makes to loop go backwards from n to 0.
            order_arg_to_sort = [arg_to_sort[arg_values_to_sort[i]] for i in range(num_individual-1,-1,-1)]
            # Replace the location values for just those features that are going to be shown individually.
            order[0:num_individual] = order_arg_to_sort[0:num_individual]
        else:
            # in descending order (most positive to minimal change to most negative)
            # sort the -1*values
            values_to_sort = [-values[order[i]] for i in range(num_individual)]
            arg_to_sort = order[0:num_individual]
            arg_values_to_sort = np.argsort(values_to_sort)
            # range(num_individual-1,-1,-1) makes to loop go backwards from n to 0.
            order_arg_to_sort = [arg_to_sort[arg_values_to_sort[i]] for i in range(num_individual-1,-1,-1)]
            # Replace the location values for just those features that are going to be shown individually.
            order[0:num_individual] = order_arg_to_sort[0:num_individual]

    # compute the locations of the individual features and plot the dashed connecting lines
    for i in range(num_individual):
        sval = values[order[i]]
        loc -= sval
        if sval >= 0:
            pos_inds.append(rng[i])
            pos_widths.append(sval)
            if lower_bounds is not None:
                pos_low.append(lower_bounds[order[i]])
                pos_high.append(upper_bounds[order[i]])
            pos_lefts.append(loc)
        else:
            neg_inds.append(rng[i])
            neg_widths.append(sval)
            if lower_bounds is not None:
                neg_low.append(lower_bounds[order[i]])
                neg_high.append(upper_bounds[order[i]])
            neg_lefts.append(loc)
        if num_individual != num_features or i + 4 < num_individual:
            plt.plot([loc, loc], [rng[i] - 1 - 0.4, rng[i] + 0.4],
                     color="#bbbbbb", linestyle="--", linewidth=0.5, zorder=-1)
        if features is None:
            yticklabels[rng[i]] = feature_names[order[i]]
        else:
            if np.issubdtype(type(features[order[i]]), np.number):
                yticklabels[rng[i]] = format_value(float(features[order[i]]), "%0.03f") + " = " + feature_names[order[i]]
            else:
                yticklabels[rng[i]] = features[order[i]] + " = " + feature_names[order[i]]

    # add a last grouped feature to represent the impact of all the features we didn't show
    if num_features < len(values):
        yticklabels[0] = "%d other features" % (len(values) - num_features + 1)
        remaining_impact = base_values - loc
        if remaining_impact < 0:
            pos_inds.append(0)
            pos_widths.append(-remaining_impact)
            pos_lefts.append(loc + remaining_impact)
            c = colors.red_rgb
        else:
            neg_inds.append(0)
            neg_widths.append(-remaining_impact)
            neg_lefts.append(loc + remaining_impact)
            c = colors.blue_rgb

    points = pos_lefts + list(np.array(pos_lefts) + np.array(pos_widths)) + neg_lefts + \
        list(np.array(neg_lefts) + np.array(neg_widths))
    dataw = np.max(points) - np.min(points)

    # draw invisible bars just for sizing the axes
    label_padding = np.array([0.1*dataw if w < 1 else 0 for w in pos_widths])
    plt.barh(pos_inds, np.array(pos_widths) + label_padding + 0.02*dataw,
             left=np.array(pos_lefts) - 0.01*dataw, color=colors.red_rgb, alpha=0)
    label_padding = np.array([-0.1*dataw if -w < 1 else 0 for w in neg_widths])
    plt.barh(neg_inds, np.array(neg_widths) + label_padding - 0.02*dataw,
             left=np.array(neg_lefts) + 0.01*dataw, color=colors.blue_rgb, alpha=0)

    # define variable we need for plotting the arrows
    head_length = 0.08
    bar_width = 0.8
    xlen = plt.xlim()[1] - plt.xlim()[0]
    fig = plt.gcf()
    ax = plt.gca()
    xticks = ax.get_xticks()
    bbox = ax.get_window_extent().transformed(fig.dpi_scale_trans.inverted())
    width, height = bbox.width, bbox.height
    bbox_to_xscale = xlen/width
    hl_scaled = bbox_to_xscale * head_length
    renderer = fig.canvas.get_renderer()

    # draw the positive arrows
    for i in range(len(pos_inds)):
        dist = pos_widths[i]
        arrow_obj = plt.arrow(
            pos_lefts[i], pos_inds[i], max(dist-hl_scaled, 0.000001), 0,
            head_length=min(dist, hl_scaled),
            color=colors.red_rgb, width=bar_width,
            head_width=bar_width
        )

        if pos_low is not None and i < len(pos_low):
            plt.errorbar(
                pos_lefts[i] + pos_widths[i], pos_inds[i],
                xerr=np.array([[pos_widths[i] - pos_low[i]], [pos_high[i] - pos_widths[i]]]),
                ecolor=colors.light_red_rgb
            )

        txt_obj = plt.text(
            pos_lefts[i] + 0.5*dist, pos_inds[i], format_value(pos_widths[i], '%+0.02f'),
            horizontalalignment='center', verticalalignment='center', color="white",
            fontsize=12
        )
        text_bbox = txt_obj.get_window_extent(renderer=renderer)
        arrow_bbox = arrow_obj.get_window_extent(renderer=renderer)

        # if the text overflows the arrow then draw it after the arrow
        if text_bbox.width > arrow_bbox.width:
            txt_obj.remove()

            txt_obj = plt.text(
                pos_lefts[i] + (5/72)*bbox_to_xscale + dist, pos_inds[i], format_value(pos_widths[i], '%+0.02f'),
                horizontalalignment='left', verticalalignment='center', color=colors.red_rgb,
                fontsize=12
            )

    # draw the negative arrows
    for i in range(len(neg_inds)):
        dist = neg_widths[i]

        arrow_obj = plt.arrow(
            neg_lefts[i], neg_inds[i], -max(-dist-hl_scaled, 0.000001), 0,
            head_length=min(-dist, hl_scaled),
            color=colors.blue_rgb, width=bar_width,
            head_width=bar_width
        )

        if neg_low is not None and i < len(neg_low):
            plt.errorbar(
                neg_lefts[i] + neg_widths[i], neg_inds[i],
                xerr=np.array([[neg_widths[i] - neg_low[i]], [neg_high[i] - neg_widths[i]]]),
                ecolor=colors.light_blue_rgb
            )

        txt_obj = plt.text(
            neg_lefts[i] + 0.5*dist, neg_inds[i], format_value(neg_widths[i], '%+0.02f'),
            horizontalalignment='center', verticalalignment='center', color="white",
            fontsize=12
        )
        text_bbox = txt_obj.get_window_extent(renderer=renderer)
        arrow_bbox = arrow_obj.get_window_extent(renderer=renderer)

        # if the text overflows the arrow then draw it after the arrow
        if text_bbox.width > arrow_bbox.width:
            txt_obj.remove()

            txt_obj = plt.text(
                neg_lefts[i] - (5/72)*bbox_to_xscale + dist, neg_inds[i], format_value(neg_widths[i], '%+0.02f'),
                horizontalalignment='right', verticalalignment='center', color=colors.blue_rgb,
                fontsize=12
            )

    # draw the y-ticks twice, once in gray and then again with just the feature names in black
    # The 1e-8 is so matplotlib 3.3 doesn't try and collapse the ticks
    ytick_pos = list(range(num_features)) + list(np.arange(num_features)+1e-8)
    plt.yticks(ytick_pos, yticklabels[:-1] + [l.split('=')[-1] for l in yticklabels[:-1]], fontsize=13)

    # put horizontal lines for each feature row
    for i in range(num_features):
        plt.axhline(i, color="#cccccc", lw=0.5, dashes=(1, 5), zorder=-1)

    if y_reverse:
        ex_loc = 1
        fx_loc = 1/num_features
    else:
        ex_loc = 1/num_features
        fx_loc = 1
    # mark the prior expected value and the model prediction with a vertical dashed grey line
    # depending on if flipping the y axis, as to which value (f(x) or E(f(x)) to put at the top or bottom)
    plt.axvline(base_values, 0, ex_loc, color="#bbbbbb", linestyle="--", linewidth=0.5, zorder=-1)
    fx = base_values + values.sum()
    plt.axvline(fx, 0, fx_loc, color="#bbbbbb", linestyle="--", linewidth=0.5, zorder=-1)

    # clean up the main axis
    plt.gca().xaxis.set_ticks_position('bottom')
    plt.gca().yaxis.set_ticks_position('none')
    plt.gca().spines['right'].set_visible(False)
    plt.gca().spines['top'].set_visible(False)
    plt.gca().spines['left'].set_visible(False)
    ax.tick_params(labelsize=13)
    #plt.xlabel("\nModel output", fontsize=12)

    # Set up the two dictionaries depending on if f(x) or E(f(x)) is at top or bottom.
    # ax2_dict contains the details for the top, ax3_dict contains the details for the bottom.
    # The code used to place these two labels uses ax.twiny() which takes the opposite position of the given axis.
    # In order to switch the location of f(x) and E(f(x)) (by setting y_reverse=True) the code swaps the order 
    # that their respective axis is created.
    if y_reverse:
        # bottom of y axis: f(x)
        ax2_dict = {"value": base_values + values.sum(),
                    "ticklabels": ["$f(x)$", "$ = "+format_value(fx, "%0.03f")+"$"],        
                    "loc1": -10/72., "loc2": 12/72., "loc3": -15/72., "loc4": -15/72.}
        # top of y axis: E(f(x))
        ax3_dict = {"value": base_values, 
                    "ticklabels": ["\n$E[f(X)]$", "\n$ = "+format_value(base_values, "%0.03f")+"$"], 
                    "loc1": -20/72., "loc2": 22/72., "loc3": 0, "loc4": 0}
        
    else:
        # bottom of y axis: E(f(x))
        ax2_dict = {"value": base_values, 
                    "ticklabels": ["\n$E[f(X)]$", "\n$ = "+format_value(base_values, "%0.03f")+"$"], 
                    "loc1": -20/72., "loc2": 22/72., "loc3": 0, "loc4": -1/72.}
        # top of y axis: f(x)
        ax3_dict = {"value": base_values + values.sum(), 
                    "ticklabels": ["$f(x)$", "$ = "+format_value(fx, "%0.03f")+"$"], 
                    "loc1": -10/72., "loc2": 12/72., "loc3": 0, "loc4": 0}

    # draw the bottom tick mark (either f(X) or E(f(x)) depending on y_reverse)
    xmin, xmax = ax.get_xlim()
    ax2 = ax.twiny()
    ax2.set_xlim(xmin, xmax)
    # The 1e-8 is so matplotlib 3.3 doesn't try and collapse the ticks
    # provide two locations as using two tick labels (one for "f(x)" and one for "= value")
    ax2.set_xticks([ax2_dict["value"], ax2_dict["value"] + 1e-8])
    ax2.set_xticklabels(ax2_dict["ticklabels"], fontsize=12, ha="left")
    ax2.spines['right'].set_visible(False)
    ax2.spines['top'].set_visible(False)
    ax2.spines['left'].set_visible(False)

    # draw the top tick mark (either f(X) or E(f(x)) depending on y_reverse)
    ax3 = ax2.twiny()
    ax3.set_xlim(xmin, xmax)
    # The 1e-8 is so matplotlib 3.3 doesn't try and collapse the ticks
    # provide two locations as using two tick labels (one for "f(x)" and one for "= value"
    ax3.set_xticks([ax3_dict["value"], ax3_dict["value"] + 1e-8])  
    ax3.set_xticklabels(ax3_dict["ticklabels"], fontsize=12, ha="left")
    # adjust the position of the bottom label (either f(X) or E(f(x)) depending on y_reverse)
    tick_labels = ax3.xaxis.get_majorticklabels()
    tick_labels[0].set_transform(tick_labels[0].get_transform(
    ) + matplotlib.transforms.ScaledTranslation(ax3_dict["loc1"], ax3_dict["loc3"], fig.dpi_scale_trans))
    tick_labels[1].set_transform(tick_labels[1].get_transform(
    ) + matplotlib.transforms.ScaledTranslation(ax3_dict["loc2"], ax3_dict["loc4"], fig.dpi_scale_trans))
    tick_labels[1].set_color("#999999")
    ax3.spines['right'].set_visible(False)
    ax3.spines['top'].set_visible(False)
    ax3.spines['left'].set_visible(False)

    # adjust the position of the top label (either f(X) or E(f(x)) depending on y_reverse)
    tick_labels = ax2.xaxis.get_majorticklabels()
    tick_labels[0].set_transform(tick_labels[0].get_transform(
    ) + matplotlib.transforms.ScaledTranslation(ax2_dict["loc1"], ax2_dict["loc3"], fig.dpi_scale_trans))
    tick_labels[1].set_transform(tick_labels[1].get_transform(
    ) + matplotlib.transforms.ScaledTranslation(ax2_dict["loc2"], ax2_dict["loc4"], fig.dpi_scale_trans))
    tick_labels[1].set_color("#999999")

    # color the y tick labels that have the feature values as gray
    # (these fall behind the black ones with just the feature name)
    tick_labels = ax.yaxis.get_majorticklabels()
    for i in range(num_features):
        tick_labels[i].set_color("#999999")
            
    if y_reverse:
        # Reverse y axis (expected at top, predicted at bottom)
        ylim = ax.get_ylim() # Get current min and max
        ax.set_ylim(ylim[1], ylim[0]) # Reverse min and max
        
    if show:
        plt.show()
    else:
        return plt.gcf()