import pandas as pd
from matplotlib import pyplot as plt
import seaborn as sns
import matplotlib.patches as mpatches
sns.set(style="whitegrid", font_scale=1.8)

def diverging_stacked_bar(results, category_names=[], colors=[], symmetrical=False,
                          bar_width = 0.35):
    if len(category_names) % 2 == 0:  # even
        print(len(category_names))
        mid = False
    else:  # odd
        mid = True
        mid_ix = len(category_names) / 2

    ind = range(len(results))
    

    fig, ax = plt.subplots(figsize=(12, 5))

    patches = []
    max_width = 0
    max_left_width = 0
    max_right_width = 0
    for i, (ix, row) in enumerate(results.iterrows()):
        b = i - bar_width / 2. # bottom_bar_edge
        if mid:
            mid_count = row.loc[category_names[mid_ix]]
            l = -1. * mid_count / 2  # left_bar_edge

            r = mpatches.Rectangle(
                    xy=(l, b), width=mid_count, height=bar_width,
                    facecolor=colors[mid_ix],
                    edgecolor=colors[mid_ix],
                    linewidth=1,
                    label='_nolegend_',
                    )
            ax.add_patch(r)
            left_edge = l
            right_edge = -l
        else:
            left_edge = right_edge = 0

        if mid:
            right_start = mid_ix + 1
            left_end = mid_ix
        else:
            right_start = len(category_names) / 2
            left_end = right_start

        # left bars
        for bar in range(0, left_end):
            bar_count = row.loc[category_names[bar]]
            l = left_edge - bar_count
            r = mpatches.Rectangle(
                    xy=(l, b), width=bar_count, height=bar_width,
                    facecolor=colors[bar],
                    edgecolor=colors[bar],
                    linewidth=1,
                    label='_nolegend_',)
            ax.add_patch(r)
            left_edge = l

        # right bars
        for bar in range(right_start, len(category_names)):
            bar_count = row.loc[category_names[bar]]
            l = right_edge
            r = mpatches.Rectangle(
                    xy=(l, b), width=bar_count, height=bar_width,
                    facecolor=colors[bar],
                    edgecolor=colors[bar],
                    linewidth=1,
                    label='_nolegend_',)
            ax.add_patch(r)
            right_edge += bar_count

        max_width = max([max_width, -left_edge, right_edge])
        max_left_width = min([max_left_width, left_edge])
        max_right_width = max([max_right_width, right_edge])

    if symmetrical:
        ax.set_xlim([-max_width, max_width])
    else:
        ax.set_xlim([max_left_width, max_right_width])
    ax.set_ylim([-bar_width, len(results)])
    ax.set_yticks([i for i in range(len(results))])
    ax.set_yticklabels([plt.Text(0, i - bar_width / 2., str(x)) for i, x in enumerate(results.index)])
    
    legend_items = []
    for cat, color in zip(category_names, colors):
        legend_items.append(mpatches.Patch(facecolor=color, edgecolor=color, label=cat))
    ax.legend(handles=legend_items)
    return ax


# Example

results = {"strongly disagree" : [3, 4, 1, 0, 5, 7],
           "disagree" : [10, 12, 3, 3, 3, 2],
           "neutral": [4, 4, 6, 8, 1, 0],
           "agree": [2, 1, 7, 4, 2, 5],
           "strongly agree": [7, 6, 5, 4, 3, 2]}
results = pd.DataFrame(results, index=["q%d" % i for i in range(6)])

columns_with_neutral = [u'agree', u'disagree', u'neutral', u'strongly agree', u'strongly disagree']
colors_with_neutral = ['firebrick', 'lightcoral', 'gainsboro', 'cornflowerblue', 'darkblue']

ax = diverging_stacked_bar(results, columns_with_neutral,
                           colors=colors_with_neutral, symmetrical=True)

columns_without_neutral = [u'agree', u'disagree', u'strongly agree', u'strongly disagree']
colors_without_neutral = ['firebrick', 'lightcoral', 'cornflowerblue', 'darkblue']

ax = diverging_stacked_bar(results, columns_without_neutral,
                           colors=colors_without_neutral, symmetrical=False)