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=(2.5, 6))

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

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

#4
text_offset = (-0.04, 0.04)
text_ha = ('right', 'left')

#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,
                                         orientation='horizontal')
#7
markerline.remove()

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

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

#10
ax.scatter(x=np.zeros(df.shape[0]), y=df['Year'], 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=text_offset[row_index % 2], y=row['Year'], s=label, 
            ha=text_ha[row_index % 2], va='center', 
            bbox={'edgecolor':'black', 'facecolor':'white'})  


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

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

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