import matplotlib as mpl
import matplotlib.pyplot as plt
from matplotlib.patches import Patch
import pandas as pd
import seaborn as sns
import numpy as np
from mpl_ornaments.titles import set_title_and_subtitle
from voting.apportionment import huntington_hill as apportion

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

df = pd.read_csv('data/age-distribution-usa-2022.csv', comment='#')
df['Percentage'] = 100*df['Population']/df['Population'].sum()

df['Labels'] = (df['Age_group'] + ' ' + 
                df['Percentage'].map('({:3.1f}%)'.format))

#1
cmap = mpl.colormaps['tab10'].resampled(df.shape[0])

#2
grid = np.zeros(shape=(10,10)) 

#3
df['Apportionement'] = apportion(votes=df['Population'], seats=grid.size)

#4
grid_data = []
for idx in df.index:
    grid_data.extend([idx]*df.iloc[idx]['Apportionement'])
grid = np.array(grid_data).reshape(grid.shape) 

#5
sns.heatmap(data=grid, ax=ax, linewidths=1.5, cmap=cmap, cbar=False)

#6
ax.set_axis_off()
#7
ax.set_aspect('equal')

#8
handles = [Patch(color=cmap.colors[idx]) for idx in df.index]

ax.legend(handles=handles, labels=df['Labels'].tolist(), ncols=1, 
          borderpad=1.2, bbox_to_anchor=(1.0, 0.5), loc='center left',
          frameon=False, title='Age group')

title = f'Population of the United States by age group (2022)'
subtitle = 'Source: US Census Bureau.'
set_title_and_subtitle(fig=fig, title=title, subtitle=subtitle, 
                       alignment='left', h_offset=102)

fig.savefig('charts/waffle-chart.png', dpi=300, bbox_inches='tight')