import geopandas as gpd
import matplotlib.pyplot as plt
from matplotlib.patches import Patch
from matplotlib import colormaps
import numpy as np
import pandas as pd
from mpl_ornaments.titles import set_title_and_subtitle

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

#1
gdp_levels = [20000, 25000, 30000, 35000, 40000, 50000, 60000] 
#2
cmap = colormaps['YlOrBr'].resampled(len(gdp_levels))

#3
gdp = pd.read_csv(f'data/uk-per-capita-gdp-by-itl2-area-2021.csv', 
                  comment='#')
gdp['level'] = np.digitize(x=gdp['per_capita_gdp'], bins=gdp_levels)

#4
boundaries = gpd.read_file(f'data/uk_itl2_geodata/'
                           f'ITL2_JAN_2021_UK_BUC_V2.shp')
boundaries.rename(columns={'ITL221CD':'ITL2_code'}, inplace=True)

#5
boundaries = boundaries.to_crs("ESRI:53004") 

#6
merged_df = pd.merge(left=boundaries, right=gdp, on='ITL2_code') 

#7
merged_df.plot(ax=ax, edgecolor='black', linewidth=0.3,
               column='level', cmap=cmap)

#8
labels, patches = list(), list()

#9
for i in range(len(gdp_levels)-1):
#10
    labels.append(f'{gdp_levels[i]}\u2014{gdp_levels[i+1]:d}') 
#11
    patches.append(Patch(edgecolor='lightgrey', facecolor=cmap(i), 
                         linewidth=0.3)) 

#12
labels.append(f'> {gdp_levels[i+1]}')
patches.append(Patch(edgecolor='lightgrey', facecolor=cmap(i+1), 
                     linewidth=0.3))

#13
ax.legend(handles=patches, labels=labels, loc = 'upper right',
          bbox_to_anchor=(1.4, 0.8), title='Pounds/yr', frameon=False,
          fontsize=8)

ax.axis('off')

title = 'Per capita GDP in the United Kingdom, 2021'
subtitle = f'By ITL2 areas. Source: Office for National Statistics.'
set_title_and_subtitle(fig=fig, title=title, subtitle=subtitle,
                       alignment='left', h_offset=90)
fig.savefig('charts/choroplet_map.png', bbox_inches='tight', dpi=300)