#tutorial 
import matplotlib.pyplot as plt
from matplotlib.offsetbox import OffsetImage, AnnotationBbox
from matplotlib.lines import Line2D
import pandas as pd

color_dict = {(2022,"Norway"): "#9194A3", (2004,"Norway"): "#2B314D",
              (2022,"Denmark"): "#E2AFA5", (2004,"Denmark"): "#A54836",
              (2022,"Sweden"): "#C4D6F8", (2004,"Sweden"): "#5375D4",
              }

xy_ticklabel_color, xlabel_color, grand_totals_color, grid_color, datalabels_color ='#C8C9C9',"#101628","#101628", "#C8C9C9", "#2B314D"

data = {
    "year": [2004, 2022, 2004, 2022, 2004, 2022],
    "countries" : [ "Denmark", "Denmark", "Norway", "Norway","Sweden", "Sweden",],
    "sites": [4,10,5,8,13,15]
}
df= pd.DataFrame(data)

#custom sort
sort_order_dict = {"Denmark":1, "Sweden":2, "Norway":3, 2004:5, 2022:4}
df = df.sort_values(by=['year','countries',], key=lambda x: x.map(sort_order_dict))
#Add the color based on the color dictionary
df['color'] = df.set_index(['year', 'countries']).index.map(color_dict.get)
df

fig, ax = plt.subplots(figsize=(5,5),facecolor = "#FFFFFF")

countries = df.countries.unique()

site_max = df.sites.max()
ax.plot(
    [list(range(len(countries)))]*site_max, 
    list(range(site_max)), 
    'o', 
    ms =2, 
    color= 
    'k')

for year, group in df.groupby("year", sort = False):
    for row in group.itertuples():
        y = row.sites
        x= row.countries
        ax.plot(
            [x]*y, 
            list(range(y)), 
            '-', 
            lw = 12, 
            color= row.color,
            solid_capstyle="round"
             )
        ax.plot(
            [x]*y, 
            list(range(y)), 
            'o', 
            ms =2, 
            color= "w" 
            )  
fig

  
ax.xaxis.set_ticks(countries, ['', '', '']) 
tick_labels = ax.xaxis.get_ticklabels()
for i,im in enumerate(img):
    ib = OffsetImage(im, zoom=.04)
    ib.image.axes = ax
    ab = AnnotationBbox(ib,
                    tick_labels[i].get_position(),
                    frameon=False,
                    box_alignment=(0.5, 2)
                    )
    ax.add_artist(ab)

fig

inset_ax = fig.add_axes([0.1, -0.15, 0.7, 0.1])  # [left, bottom, width, height]
inset_ax.axis('off')

colors = df.color.unique()
x0 = 0
width = .2  
y = 0

for i, color in enumerate(colors[:3]):
    inset_ax.plot([x0 + i, x0 + i + 1], [y, y], color=color, lw=4)

inset_ax.text(x0 + 3.2, y, 'Before 2004', va='center', fontsize=12)

# Shift x for second group
x0 = x0 + 9  # Add space between the two groups visually

for i, color in enumerate(colors[3:7]):
    inset_ax.plot([x0 + i, x0 + i + 1], [y, y], color=color, lw=4)

inset_ax.text(x0 + 3.2, y, 'After 2004', va='center', fontsize=12)

# Set x/y limits of inset axes
inset_ax.set_xlim(-1, x0 + 5)
inset_ax.set_ylim(-1, 1)

fig

ax.set_frame_on(False)
ax.tick_params(axis='both', which='major', length=0, labelleft = False)

fig