import geopandas as gpd
from shapely.geometry import Polygon
import matplotlib.pyplot as plt
import pandas as pd
import seaborn as sns
from mpl_ornaments.titles import set_title_and_subtitle

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

#1
bounds = Polygon(shell=[sw := (-25,30), nw := (-25,72), ne := (60,72),
                        se := (60,30)]) 
#2
iso_country_codes = pd.read_csv(
    'data/iso-country-codes-with-continents.csv', comment='#') 
iso_country_codes_eu = pd.read_csv('data/iso-country-codes-eu.csv',
                                   comment='#')

#3
gdata = gpd.read_file(f'data/world-administrative-boundaries/'
                      f'world-administrative-boundaries.shp') 
gdata = gdata[gdata['iso3'].isin(iso_country_codes_eu['iso3'])]
gdata = gdata.clip(mask=bounds)

#4
df_approx_centroids = pd.read_csv('data/eu-approximate-centroids.csv',
                                  comment='#')
df_approx_centroids.dropna(inplace=True)

#5
df_edu_spending = pd.read_csv('data/eu-education-spending-2020.csv',
                              comment='#')                        
#6
df_data_points = pd.merge(
    left=df_approx_centroids, 
    right=df_edu_spending[['iso3', 'education_spending']], on='iso3')

#7
gdata.plot(ax=ax, color='floralwhite', edgecolor='black', linewidth=0.3) 

#8
sns.scatterplot(data=df_data_points, x='lon', y='lat',
                size='education_spending', ax=ax, 
                hue='education_spending', sizes=(25,300), 
                alpha=0.7, edgecolor='dimgray')

#9
legend = ax.get_legend()
legend.set_title('% of GDP')
sns.move_legend(ax, loc='upper left', frameon=False)

ax.axis('off') 

title = 'Education spending in EU countries'
subtitle = f'Primary to post-secondary. Source: OECD (2020 or latest av.)'
set_title_and_subtitle(fig=fig, title=title, subtitle=subtitle,
                       alignment='left', h_offset=80)
fig.savefig('charts/proportional_symbol_map.png', bbox_inches='tight', 
            dpi=300)