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

#1
country_codes_g7 = ['CAN', 'DEU', 'FRA', 'GBR', 'ITA', 'JPN', 'USA']
df = pd.read_csv('data/education-spending-and-level-g20.csv', comment='#')
df = df[df['Country_code'].isin(country_codes_g7)]
df.sort_values(by='Tertiary_educated', ascending=False, inplace=True)

#2
cmap = ['#f0e998', '#9b976e']

#3
ncols = 4
nrows = np.ceil(len(country_codes_g7)/ncols).astype(int)
grid_spec = {'hspace': -0.3, 'wspace': 0.2}
fig, axes = plt.subplots(nrows=nrows, ncols=ncols, gridspec_kw=grid_spec)
axes = axes.flatten()

#4
for idx, (_, row) in enumerate(df.iterrows()):
#5
    rounded_percentage = np.round(row['Tertiary_educated']).astype(int)
#6
    grid = np.ones(shape=(10,10))
    grid.reshape(-1)[0:(100-rounded_percentage)] = 0
#7
    sns.heatmap(data=grid, ax=axes[idx], linewidths=1.5, cmap=cmap, 
                cbar=False)
#8
    axes[idx].set_title(label=(f'{row["Country_name"]} '
                               f'({row["Tertiary_educated"]:3.1f}%)'),
                        fontsize=10)

for ax in axes:
    ax.set_axis_off()
    ax.set_aspect('equal')    

title = f'Fraction of tertiary-educated people in the G7'
subtitle = 'Source: OECD Data (2020 or latest available)'
set_title_and_subtitle(fig=fig, title=title, subtitle=subtitle, 
                       alignment='left', h_offset=60, v_offset=20)
    
fig.savefig('charts/waffle-chart-multiple.png', dpi=300, 
            bbox_inches='tight')