import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from mpl_ornaments.titles import set_title_and_subtitle

fig, ax = plt.subplots(nrows=1, ncols=1, figsize=(6.5, 7.5))

#1
df = pd.read_csv('data/ai-discovery-cancer-care.csv', comment='#')
likert_levels = list(df.columns); likert_levels.remove('Application')

#2
nlevels = len(likert_levels)
likert_colors = plt.get_cmap('coolwarm')(np.linspace(0.1, 0.9, nlevels))

#3
upper_bound = np.floor(nlevels/2).astype(int)
centroid = df[likert_levels[0:upper_bound]].apply(sum, axis=1)
if (nlevels % 2):
    centroid = centroid + df[likert_levels[upper_bound]]/2

#4
left = -centroid
for likert_level, likert_color in zip(likert_levels, likert_colors):
    perceived_likelihood = df[likert_level]    
    
    ax.barh(y=df['Application'], width=perceived_likelihood, left=left, 
            label=likert_level, color=likert_color) 
    
    left = left + perceived_likelihood

#5
ax.axvline(x=0.0, color='white', linewidth=1.0) 

#6
ax.set_xlim(left = 1.05*ax.get_xlim()[0]) 

ax.spines[['right', 'top']].set_visible(False)
ax.legend(ncols=len(likert_levels), bbox_to_anchor=(0.55, 1.03), 
          loc='center', frameon=False)
   
title = ('Which AI application will be more likely succesful in cancer '
         'treatment?')
subtitle = ('Within ten years from survey. Source: Cabral et al., Curr. '
            'Oncol., 2023.')
set_title_and_subtitle(fig=fig, title=title, subtitle=subtitle,
                       alignment='left', h_offset=-20)

fig.savefig('charts/hundred-percent-divergent-stacked-bar-chart.png', 
            bbox_inches='tight', dpi=300)