import matplotlib.pyplot as plt
from matplotlib.offsetbox import OffsetImage, AnnotationBbox
import numpy as np
import pandas as pd

xy_ticklabel_color, year_label = "#101628", "#F5F7F7"

data = {
    "year": [2004, 2022, 2004, 2022, 2004, 2022],
    "countries" : ["Sweden", "Sweden", "Denmark", "Denmark", "Norway", "Norway"],
    "sites": [13,15,4,10,5,8]
}
df= pd.DataFrame(data)
df['ctry_code'] = df.countries.astype(str).str[:2].astype(str).str.upper()
df = df.sort_values([ 'year'], ascending=True )

df

img = [
    plt.imread("../flags/sw-sq.png"),
    plt.imread("../flags/de-sq.png"), 
    plt.imread("../flags/no-sq.png")] * 3

countries = df.countries.unique()

y_end = np.repeat([1,-1], len(countries))
y_start =np.repeat([0.1,-0.2], len(countries))
max_sites = df.sites.max()
print(y_start, y_end)

fig, ax = plt.subplots( figsize=(10,5))

ax.set_ylim(-1,1)

#add the numbers using the x labels
ax.xaxis.set_ticks(np.arange(0, max_sites + 1, 1), labels = np.arange(0,max_sites + 1,1) )
ax.spines['bottom'].set_position('zero')
ax.tick_params(
    axis='both', 
    which='major',
    length=0, 
    labelsize=16, 
    labelleft = False, 
    colors =xy_ticklabel_color
    )
ax.set_frame_on(False)

for site, im, ys, ye in zip(df.sites, img, y_start, y_end):
    image_box =  OffsetImage(im, zoom = 0.04) #container for the image
    ab = AnnotationBbox(
        image_box, 
        (site, ye), 
        frameon = False
        )
    ax.add_artist(ab)

    ax.annotate(
        "", 
        (site, ye),
        (site , ys),
        ha="right", 
        va="center",
        zorder=1,
        arrowprops=dict(
            arrowstyle="]-",
            fc="w", 
            ec="#D8DDDF",
            connectionstyle="arc3,rad=-0")
        )

fig

pos = [
    (max_sites / 2,0.5),
    (max_sites /2,-1.5)
    ]
    
for pos, year in zip(pos, df.year.unique()):
    ax.annotate(
        year, 
        xy = pos, 
        size = 150, 
        ha = "center", 
        color = year_label, 
        weight = "bold", 
        annotation_clip = False, 
        zorder = 0)

fig