from collections import namedtuple
from matplotlib.patches import Patch
import matplotlib.pyplot as plt
import pandas as pd
from mpl_ornaments.titles import set_title_and_subtitle

fig, ax = plt.subplots(nrows=1, ncols=2, figsize=(6.4,5.5))
ax = ax.flatten()

df = pd.read_csv('data/projected-age-distribution-uk-2022.csv',
                 comment='#')

#1
Gender = namedtuple("Gender", ["name", "colour"])
genders = [Gender('Male', '#80b1d3'), Gender('Female', '#fb8072')]

#2
df_wide = df.pivot(index='Age_group', columns='Gender', 
                   values='Projected_pop')

#3
yticklabels = df_wide.index.values
yticks = list(range(len(yticklabels)))

for idx, gender in enumerate(genders):
#4
    ax[idx].barh(y=yticks, width=df_wide[gender.name].values, 
                 color=gender.colour)    
    
    ax[idx].spines[['right', 'top', 'left']].set_visible(False)

#5
ax[0].invert_xaxis()
ax[0].set_yticks([])

#6
ax[0].legend(handles=[Patch(color=genders[0].colour)], 
             labels=[genders[0].name], loc='upper left', frameon=False)

#7
ax[1].set_yticks(yticks)
#8
ax[1].tick_params(axis='y', length=0)

#9
ax[1].set_yticklabels(yticklabels)

#10
ax[1].legend(handles=[Patch(color=genders[1].colour)], 
             labels=[genders[1].name], loc='upper right', frameon=False)

title = 'Age distribution in the United Kingdom (2022, projected data)'
subtitle = ('In thousand persons. Source: UK Office for National '
            'Statistics.')
set_title_and_subtitle(fig=fig, title=title, subtitle=subtitle,
                       alignment='left', h_offset=20)

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