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

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

xy_ticklabel_color, grand_totals_color, grid_color, datalabels_color, title_color ='#101628',"#101628", "#C8C9C9", "#FFFFFF", "#7C8091"

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

df['ctry_code'] = df.countries.map(code_dict)

#custom sort
sort_order_dict = {"Denmark":2, "Sweden":1, "Norway":3, 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, axes = plt.subplots(ncols=len(df.year.unique()), nrows=1, figsize=(7,5), sharex=True, sharey=True, facecolor="#FFFFFF", zorder=1)
plt.subplots_adjust(wspace=1.7) 
fig.tight_layout() 

for i, (year, group) in enumerate(df.groupby("year", sort = False)):
    # Sweden
    sweden_values = group[group['countries'] == 'Sweden']['sites'].values[0]
    sweden_code = group[group['countries'] == 'Sweden']['ctry_code'].values[0]

    sweden_bar = axes[i].bar(sweden_code, sweden_values, color=color_dict["Sweden"], width=0.4)
    axes[i].bar_label(sweden_bar, label_type='edge', weight = "bold") 
    axes[i].text(sweden_code, 0.5, sweden_values, ha = "center", color = "w", weight = "bold")
    axes[i].text(sweden_code, 1.5, sweden_code, ha = "center", color = "w")

    # Denmark + Norway 
    denmark_val = group[group['countries'] == 'Denmark']['sites'].values[0]
    norway_val = group[group['countries'] == 'Norway']['sites'].values[0]
    sk_no = "DK + NO"

    axes[i].bar(sk_no, norway_val, color=color_dict["Norway"], width=0.4, label='Denmark')
    denmark_bar = axes[i].bar(sk_no, denmark_val, bottom=norway_val, color=color_dict["Denmark"], width=0.4, label='Norway')
    axes[i].bar_label(denmark_bar, label_type='edge', weight = "bold")  # Add Denmark label
    axes[i].text(sk_no, 0.5, norway_val, ha = "center", color = "w", weight = "bold")
    axes[i].text(sk_no, 0.5 + norway_val, denmark_val, ha = "center", color = "w", weight = "bold")
    axes[i].text(sk_no, 1.5, group[group['countries'] == 'Norway']['ctry_code'].values[0], ha = "center", color = "w", )
    axes[i].text(sk_no, 1.5 + norway_val, group[group['countries'] == 'Denmark']['ctry_code'].values[0], ha = "center", color = "w", )
    
    axes[i].set_title(str(year), color = grand_totals_color)
    axes[i].set_frame_on(False)
    axes[i].tick_params(axis='both', which='major', length=0, labelleft = False, labelsize=9, colors= xy_ticklabel_color, pad =50)

    tick_positions = [sweden_code, sk_no]  

    # Add flag images at tick positions
    for pos, im in zip(tick_positions, img):
        imagebox = OffsetImage(im, zoom=0.04)
        ab = AnnotationBbox(imagebox, (pos, 0), 
                          frameon=False, 
                          box_alignment=(0.5, 2),  # Adjust alignment as needed
                          pad=0)
        axes[i].add_artist(ab)

        image_box =  OffsetImage(no, zoom = 0.04) #container for the image
        ab = AnnotationBbox(image_box, (sk_no,0), xybox= (1.1, -1.8),  frameon = False)
        axes[i].add_artist(ab)

#get positions of the subplots and calculate the middle
pos1 = axes[0].get_position()
pos2 = axes[1].get_position()

middle_x = (pos1.x1 + pos2.x0) / 2

# Add line spanning full height of figure
line = plt.Line2D([middle_x, middle_x], [0, 1], 
                 color="#E1E5E7", linewidth=1, 
                 transform=fig.transFigure)
fig.add_artist(line)

fig