import matplotlib.pyplot as plt
import numpy as np
import pandas as pd

color_dict = {"Norway": "#2B314D", "Denmark": "#A54836", "Sweden": "#5375D4" }

xy_ticklabel_color, grid_color, datalabels_color ="#101628", "#D7DCDE", "#FFFFFF"

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['pct_tot'] = df['sites'] / df.groupby('year')['sites'].transform('sum')
df['sub_total'] = df.groupby('year')['sites'].transform('sum')
df['xticks'] = df['sub_total'] / 2
#custom sort
sort_order_dict = {"Denmark":2, "Sweden":3, "Norway":1, 2004:4, 2022:5}
df = df.sort_values(by=['countries','year',], key=lambda x: x.map(sort_order_dict))

#map the colors of a dict to a dataframe
df['color']= df.countries.map(color_dict)
df

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

bottom = np.zeros(len(df.year.unique()))
for country, group  in df.groupby("countries", sort = False):

    #get the  x position
    total = df.sub_total.unique()
    cumsum = np.cumsum(np.insert(total[:-1], 0, 0))
    xticks = cumsum + total / 2
    
    y = group.pct_tot.tolist()

    ax.bar(
        xticks, 
        height = y, 
        width= total, 
        bottom = bottom, 
        align='edge', 
        color=group.color, 
        )
    bottom +=y

#add data labels
for bar, row in zip(ax.patches,df.itertuples()):
    x= bar.get_x() + bar.get_width() / 2

    ax.text(
        bar.get_x() + bar.get_width() / 2,            
        bar.get_height()/2 + bar.get_y() ,
        f'{row.pct_tot:.1%}',                          
        ha='center', 
        va="center",
        color=datalabels_color,  
        size=12 
        )
    


fig

y_start = -0.1
y_end = -0.05

for i, (year, group) in enumerate(df.groupby("year")):
    sub_tot = group.sub_total.iloc[0]
    center_of_bar = xticks[i] + sub_tot /2 

    ax.text(
        center_of_bar, 
        -0.2 , 
        f'{sub_tot}\n{group.year.iloc[0]}', 
        ha='center', 
        size = 12, 
        color = xy_ticklabel_color
        ) 


    ax.annotate(
        '', 
        xy=(center_of_bar, y_end), #destination point
        xytext=(center_of_bar, y_start), #start point
        annotation_clip=False, 
        xycoords='data', 
        arrowprops=dict(
            arrowstyle=f'-[, widthB={sub_tot *0.3}, lengthB=0.1', 
            lw=1.0,
            color=grid_color)
            )


ax.set_axis_off()
fig