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)

df = pd.read_csv('data/top-10-trading-partners-usa-2017.csv',
                 comment='#')
df['Percentage'] = 100*df['Total_trade']/df['Total_trade'].sum()
    
df['labels'] = (df['Country'] + ' ' + 
                df['Percentage'].map('{:3.1f}%'.format))
        
rectprops = {'ec':'w'}
textprops = {'c':'k', 'reflow':True, 'max_fontsize':14}

#1
tmap = tr.treemap(axes=ax, data=df, area='Percentage', 
                  labels=df['labels'].to_list(), fill='Continent', 
                  cmap='Set3', textprops=textprops, 
                  rectprops=rectprops, levels=['Continent', 'Country'])

#2
continents = sorted(df['Continent'].unique().tolist())
handles = tmap.handles['Country']
handles.sort(key=lambda x: x.get_label())

#3
legend = fig.legend(handles=handles, labels=continents, 
                    loc='lower center', ncol=len(continents), 
                    frameon=False)

ax.set_axis_off()

title = (f'Top ten United States trade partners, 2017')
subtitle = (f'Percentage of total trade in goods. Source: '
            f'US Census Bureau.')
set_title_and_subtitle(fig=fig, title=title, subtitle=subtitle,
                       alignment='left', h_offset=60)

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