import geoplot as gplt
import geoplot.crs as gcrs
import geopandas as gpd

import pandas as pd
import seaborn as sns
from mpl_ornaments.titles import set_title_and_subtitle

#1
gdata = gpd.read_file(f'data/'
                      f'cb_2018_us_state_500k/cb_2018_us_state_500k.shp')

#2
us_states = pd.read_csv(f'data/us-states.csv', comment='#',
                        dtype={'contiguous': 'bool'})
us_states = us_states[us_states['contiguous']]

#3
us_gdp_by_state = pd.read_csv(f'data/us-gdp-by-state-2020.csv', 
                              comment='#')
us_gdp_by_state['gdp'] = us_gdp_by_state['gdp']/1000 

#4
gdata = pd.merge(left=gdata, right=us_states, left_on='NAME',
                 right_on='name')
gdata = pd.merge(left=gdata, right=us_gdp_by_state, left_on='NAME',
                 right_on='name')

#5
ax = gplt.cartogram(
    gdata, scale='gdp', hue='gdp', projection=gcrs.AlbersEqualArea(),
    legend=True, legend_var='hue', limits=(0.6,0.9), cmap='summer_r', 
    edgecolor='dimgrey', linewidth=1, 
    legend_kwargs={'shrink':.5, 'location':'top'},
)

#6
gplt.polyplot(gdata, facecolor='lightgray', edgecolor='white', ax=ax)

ax.axis('off') 
fig = ax.get_figure()

title = 'GDP by US state (2022)'
subtitle = f'In bn $. Source: US Bureau of Economic Analysis'
set_title_and_subtitle(fig=fig, title=title, subtitle=subtitle,
                       alignment='left', h_offset=80)
fig.savefig('charts/cartogram.png', bbox_inches='tight', dpi=300)