import matplotlib.pyplot as plt
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)

df = pd.read_csv('data/eu-per-capita-gdp-by-country-2021.csv',
                 comment='#')
df['Per_capita_gdp'] = df['Per_capita_gdp']/10**3

bin_width, lbound, ubound = 15.0, 0.0, 135.0
bin_edges = np.arange(start=lbound, stop=ubound, step=bin_width)
bin_edges = np.append(bin_edges, ubound)

#1
counts, _ = np.histogram(a=df['Per_capita_gdp'], bins=bin_edges)
centroids = (bin_edges[:-1] + bin_edges[1:])/2 

#2
for count, centroid in zip(counts, centroids):
    ax.scatter(x=[centroid]*count, y=np.arange(start=1, stop=count+1), 
               marker='o', s=80, color='#1f77b4')

#3
xtickslabels = [f'{bin_edges[i]:.0f}-{bin_edges[i+1]:.0f}' 
                for i, _ in enumerate(bin_edges[:-1])]

ax.set_xticks(ticks=centroids)
ax.set_xticklabels(labels=xtickslabels)

ax.spines[['right', 'top', 'left']].set_visible(False)

title = 'Distribution of per capita GDP in the EU (2021)'
subtitle = 'In current k USD; one dot = one country. Source: The World Bank.'
set_title_and_subtitle(fig=fig, title=title, subtitle=subtitle,
                       alignment='left', h_offset=40)

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