import geopandas as gpd
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from matplotlib.lines import Line2D

color_dict = {2022: "#5375D4", 2004: "#CC5A43", }

xy_ticklabel_color,color_map ='#101628','#D3D3D3'

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['diff']=df.groupby('countries')['sites'].diff().fillna(df.sites).astype(int)

df['sub_total'] = df.groupby('countries')['diff'].transform('sum')
df = df.sort_values([ 'year', 'sub_total'], ascending=True )
df['ctry_code'] = df.countries.astype(str).str[:2].astype(str).str.upper()
df['color']= df.year.map(color_dict)
df

url = "https://raw.githubusercontent.com/eurostat/Nuts2json/master/pub/v2/2021/3035/20M/0.json"
map_df = gpd.read_file(url,  layer='nutsrg')
map_df = map_df[map_df.id.isin(['NO','SE', 'DK']) ]

map_df

fig = plt.figure(figsize=(15,10))
ax_map = fig.add_axes([0, 0, 1, 1])
map_df.plot(color=color_map,ax=ax_map)
ax_map.set_axis_off()

for i, (country, group) in enumerate(df.groupby("countries", sort = False)):
    axes_height = 0.4
    axes_width = 0.03
    #add the axes at lat and lon
    ax_bar = fig.add_axes([lat[i] ,lon[i] , axes_width, axes_height]) 

    bottom = np.zeros(1)
    
    for row in group.itertuples():

        y = row.diff
        ax_bar.bar(
            range(1), 
            y, 
            bottom = bottom, 
            width= 0.2, 
            color=row.color
            )
        bottom +=y

        #add grand totals
        ax_bar.text(
            0, 
            row.sub_total+1, 
            row.sub_total, 
            ha= "center", 
            color = "white",
                bbox=dict(
                    facecolor='black', 
                    edgecolor='black', 
                    boxstyle='round,pad=0.6'
                    )
            )
        #add the x labels
        ax_bar.set_xlabel(
            row.ctry_code, 
            color =xy_ticklabel_color, 
            size =14, 
            )

        #add the data labels
        offset_text = 0.3
        for bar in ax_bar.patches:
            ax_bar.text(
                bar.get_x() + bar.get_width() / 2,            
                bar.get_height()/2 + bar.get_y(),                
                round(bar.get_height()),                          
                ha='center', 
                color='w',  
                size=10
                )

    #add styling
    ax_bar.set_ylim(0,15)
    ax_bar.set_frame_on(False)
    ax_bar.tick_params(
        axis='both', 
        which='both',
        length = 0, 
        labelleft = False, 
        labelbottom = False
        )

fig

labels = ["Before 2004", " After 2004", "Total"]
colors = ["#CC5A43","#5375D4", "black"]

lines = [Line2D([0], [0], color=c, linestyle='-', ) for c in colors]

spacing = 0.08
y_line = 0.01
# Create legend manually
for i, (line, label) in enumerate(zip(lines, labels)):
    # Add label as text above
    fig.text(
        0.62 + i * spacing, 
        0, 
        label, 
        ha='center', 
        va='bottom', 
        fontsize=12, 
        clip_on=False,
        color=xy_ticklabel_color
        )

    # Add line below the label
    fig.lines.append(
        Line2D(
            [0.6 + i * spacing, 0.65 + i * spacing],  # x-range
            [-y_line, -y_line],  # y-position (below text)
            color=line.get_color(),
            lw=4,
            transform=fig.transFigure,
            clip_on=False,
        )
    )
fig