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

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

#1
df = pd.read_csv('data/education-spending-and-level-g20.csv', comment='#')
x, y, z, txt = df['Rel_spending'], df['Tertiary_educated'],\
               df['Abs_spending'], df['Country_code']

#2
sns.scatterplot(data=df, x=x, y=y, size=z, sizes=(100,1000), alpha=0.4,
                ax=ax)

#3
handles, labels = ax.get_legend_handles_labels()

#4
for x_, y_, txt in zip(x, y, df['Country_code']):
    ax.annotate(txt, (x_, y_), ha='center', va='center')

#5
labels = [f'{float(l)/1000:3.1f}' for l in labels]

#6
for h in handles[-6::]:
    h.set_color('lightgrey'); h.set_markeredgecolor('grey')

#7
leg = ax.legend(handles=handles, labels=labels, borderpad=1.5, 
                labelspacing=1.0, ncols=2, handletextpad=1.0)

#8
leg.set_title('Total spending, tertiary\n[k$ student/year]', 
              prop={'weight':'bold'})
leg.get_title().set_multialignment('center')

#9
ax.set_xlim(xmin=0.85*min(x), xmax=1.15*max(x))
ax.set_ylim(ymin=0.80*min(y), ymax=1.05*max(y))

#10
ax.grid(which='major', axis='both')
ax.set_xlabel('Education spending, tertiary [% of GDP]')
ax.set_ylabel('Adult education level, tertiary [% of pop.]')
ax.spines[['right', 'top']].set_visible(False)

title = 'Education spending vs. adult education level (tertiary)'
subtitle = ('Selected G20 countries, 2020 or latest available data. '
            'Source: OECD Data.')
set_title_and_subtitle(fig=fig, title=title, alignment='left',
                       subtitle = subtitle, h_offset = 20)

fig.savefig('charts/scatter-bubble-plot-with-annotations.png', 
            bbox_inches='tight', dpi=300)