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

#1
fig, ax = plt.subplots(figsize=(11, 2.5), constrained_layout=True)

#2
df = pd.read_csv('data/british-pm.csv', comment='#')
df = df[df['Year'].between(left=1970, right=2010)]

#3
cpalette = colormaps['Dark2'].colors
extensions_col, baseline_col, stems_and_bullets_col = cpalette[0:3]

#4
text_offset = (-0.3, 0.3)
text_va = ('top', 'bottom')

#5
stem_tips = np.zeros(df.shape[0])
stem_tips[::2] = text_offset[0]; 
stem_tips[1::2] = text_offset[1]   

#6
markerline, stemline, baseline = ax.stem(df['Year'], stem_tips)
#7
markerline.remove()

#8
plt.setp(stemline, color=stems_and_bullets_col, zorder=0)
plt.setp(baseline, color=baseline_col, zorder=0)

#9
delta = 2
start_yr, end_yr = df['Year'].apply([min, max])
ax.plot((start_yr-delta, start_yr), (0,0), c=extensions_col, zorder=0)
ax.plot((end_yr, end_yr+delta), (0,0), c=extensions_col, zorder=0)

#10
ax.scatter(x=df['Year'], y=np.zeros(df.shape[0]), s=120, 
           color=stems_and_bullets_col)

#11
for row_index, row in df.iterrows():

#12
    label = (f'{row["Year"]:d}\n'
             f'{row["Name"]} {row["Surname"]}\n'
             f'({row["Party"]})')
    
#13
    ax.text(x=row['Year'], y=text_offset[row_index % 2], 
            va=text_va[row_index % 2], s=label, ha='center')  

#14
ax.spines[["left", "top", "right", "bottom"]].set_visible(False)
ax.set_xticks([]); ax.set_yticks([])

#15
ax.set_ylim(-1.5, 1.5)

title = (f'British Prime Ministers between {start_yr} and {end_yr}')
set_title_and_subtitle(fig=fig, title=title, subtitle=None,
                       alignment='center')


fig.savefig('charts/horizontal-timeline.png', bbox_inches='tight',
            dpi=400)