import matplotlib.pyplot as plt
from matplotlib.lines import Line2D
import pandas as pd

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

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)

sort_order_dict = {"Denmark": 2, "Sweden": 1, "Norway": 3, 2004: 5, 2022: 4}
df = df.sort_values(by=["year","countries"], key=lambda x: x.map(sort_order_dict))

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

fig, axes = plt.subplots(nrows=df.countries.nunique(),figsize=(10,8), sharex=True)

for ax, (country, group) in zip(axes, df.groupby("countries", sort=False)):
    num_bars = group[group['year']==2022].sites
    color_count = group[group['year']==2004].sites
    primary_color =  group[group['year']==2022].color.iloc[0] # 
    secondary_color = group[group['year']==2004].color.iloc[0]  

    for num_bars, color_count in zip(num_bars, color_count):

        full_groups = num_bars // 5
        remainder = num_bars % 5
        #print(full_groups, remainder, num_bars, color_count, primary_color)
        
        bars_colored = 0
        
        # plot complete groups (5 bars each)
        for group in range(full_groups):
            offset = group * 5
            for i, x in enumerate([1 + offset, 2 + offset, 3 + offset, 4 + offset]):
                color = primary_color if bars_colored < color_count else secondary_color
                ax.plot([x, x], [0, 20], '-', linewidth=12, color=color, solid_capstyle="round")
                bars_colored += 1
            
            #diagonal bar
            color = primary_color if bars_colored < color_count else secondary_color
            ax.plot([1 + offset, 4 + offset], [5, 15], '-', linewidth=12, color=color, solid_capstyle="round")
            bars_colored += 1
        
        #plot remaining bars
        if remainder > 0:
            offset = full_groups * 5
            vertical_count = min(remainder, 4)
            
            for x in range(1 + offset, offset + vertical_count + 1):
                color = primary_color if bars_colored < color_count else secondary_color
                ax.plot([x, x], [0, 20], '-', linewidth=12, color=color, solid_capstyle="round")
                bars_colored += 1
            
            # diagonal bar only if we have the full pattern (4 vertical + 1 diagonal)
            if remainder == 5:
                color = primary_color if bars_colored < color_count else secondary_color
                ax.plot([1 + offset, 4 + offset], [5, 15], '-', linewidth=12, color=color, solid_capstyle="round")
                bars_colored += 1

        ax.set_ylim(-5, 25)
        ax.set_ylabel(country, size = 14,rotation=0, labelpad = 30)
        ax.tick_params(axis='both', which='major', length=0, labelleft= False, labelbottom=False, )
        ax.set_frame_on(False)

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)

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)

inset_ax.set_xlim(-1, x0 + 5)
inset_ax.set_ylim(-1, 1)

fig