import matplotlib.pyplot as plt
import pandas as pd
import mpl_extra.treemap as tr
from mpl_ornaments.titles import set_title_and_subtitle

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

#1
df = pd.read_csv('data/top-10-causes-of-death-usa-2021.csv', comment='#')
df['Percentage'] = 100*df['Number_of_fatalities']/\
    df['Number_of_fatalities'].sum()

#2
df['labels'] = (df['Cause_of_death'] + ' ' + 
                df['Percentage'].map('{:3.1f}%'.format))
    
#3
rectprops = {'ec':'w'}
textprops = {'c':'k', 'reflow':True, 'max_fontsize':14}

#4
tr.treemap(axes=ax, data=df, area='Percentage', 
           labels=df['labels'].to_list(), fill='Cause_of_death',
           cmap='Set3', textprops=textprops, rectprops=rectprops)

#5
ax.set_axis_off()

title = ('The ten leading causes of death in the US, 2021')
subtitle = ('Source:  CDC/National Center for Health Statistics.')
set_title_and_subtitle(fig=fig, title=title, subtitle=subtitle,
                       alignment='left', h_offset=60)

fig.savefig('charts/one-level-tree-map.png', dpi=300, bbox_inches='tight')