import matplotlib.pyplot as plt

import pandas as pd
import squarify

from matplotlib.patches import FancyBboxPatch

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

xy_ticklabel_color, xlabel_color,  datalabels_color ='#757C85',"#101628",  "#FFFFFF"

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)
df['sub_total'] = df.groupby('year')['sites'].transform('sum')
df['ctry_code'] = df.countries.astype(str).str[:2].astype(str).str.upper()
#map the colors of a dict to a dataframe
df['color']= df.countries.map(color_dict)

df

fig, axes = plt.subplots(ncols = df.year.nunique(), nrows = 1, sharex = True, figsize=(10,5), facecolor = "#FFFFFF", zorder= 1)
fig.tight_layout(pad=1.0)

scalexy = [75,100]
y_title = [-1, 0]

for i, ((year, group), ax) in enumerate(zip(df.groupby("year"), axes.ravel())):
    # Normalize sizes to fit within a 100x100 square
    normed_sizes = squarify.normalize_sizes(group['sites'].tolist(), scalexy[i], scalexy[i])

    # Get squarified rectangles but plot them only with fancybboxpatch
    rects = squarify.squarify(normed_sizes, 0, 0, scalexy[i], scalexy[i])

    for j, (rect, row) in enumerate(zip(rects, group.itertuples())):
        x, y, dx, dy = rect['x'], rect['y'], rect['dx'], rect['dy']

        # Draw rounded rectangle
        patch = FancyBboxPatch(
            (x, y), dx, dy,
            boxstyle="round,pad=0.02,rounding_size=2",
            linewidth=1,
            facecolor=row.color,
            edgecolor='w'
        )
        ax.add_patch(patch)

        #add the data labels
        ax.text(
            x + 5,
            y + dy -5,
            f'{row.ctry_code}\n{row.sites}',
            ha='center',
            va='center',
            fontsize=8,
            color='w'
        )
    
        if j == 1:

            #add the totals
            ax.text(
                x  ,
                y + dy + 2,   
                f' Total {row.sub_total}', 
                size = 8,
                color=xy_ticklabel_color
                )

    #add year labels
    ax.set_xlabel(
        year, 
        color=xlabel_color, 
        size=14, 
        weight="bold", 
        labelpad = 53
        )
    
    # Styling
    ax.set_ylim(0, 100)
    ax.set_xlim(0, scalexy[i])
    ax.set_aspect('equal')
    ax.set_frame_on(False)
    ax.tick_params(
        axis='both', 
        which='both',
        length = 0, 
        labelleft = False,
        labelbottom = False
        )