import math
import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1.inset_locator import inset_axes
import numpy as np
import pandas as pd
import geopandas as gpd



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


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

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

#custom sort
sort_order_dict = {"Denmark":3, "Sweden":1, "Norway":2, 2004:5, 2022:4}
df = df.sort_values(by=['countries','year',], key=lambda x: x.map(sort_order_dict))
df['ctry_code'] = df.countries.map(code_dict)
df['color']= df.countries.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, axes = plt.subplots(nrows=1, ncols=3, figsize=(10,3),)
fig.subplots_adjust(wspace=1)  # vertical spacing

for ax, (country, group) in zip(axes.ravel(), df.groupby("ctry_code", sort = False)):
    
    #add the maps
    ctry_map = map_df[map_df['id'] == country].plot(ax=ax, color = color_map)
    ax.set_axis_off()
    
    # create an overlay axes (inset) in the same position 
    ax_matrix = inset_axes(ax, width=1, height=1, loc='center')
    ax_matrix.invert_xaxis()
    edge_color = np.insert( 
        np.repeat(
            group.color, group['diff']), 
            0,  
            ["#FFFFFF"] * (matrix_columns * matrix_rows - group.sub_total.iloc[0]
            ))


    face_color = np.insert( np.repeat(
        ["w", group.color.iloc[0]], group['diff']), 
        0,  
        ["#FFFFFF"] * (matrix_columns * matrix_rows - group.sub_total.iloc[0]
        ))

    color = group.color.iloc[0]
    #add the matrix
    ax_matrix.scatter(
        y, 
        x, 
        s= 120,
        color = face_color,
        ec = edge_color,
        clip_on = False
        )

    #add the country labels
    ax_matrix.set_title(
        group.countries.iloc[0], 
        y = 2,
        size = 8,
        color= xy_ticklabel_color
        )

    #add the site labels
    ax_matrix.set_xlabel(
        f'{group.sites.iloc[1]}', 
        color =xy_ticklabel_color, 
        size =26,
        weight = "bold",
        labelpad = 70
        )

    #add the site count
    ax_matrix.annotate(
        f'+{int(group['diff'].iloc[0])} ', 
        xy =(3.5,-4.5), 
        color = color, 
        annotation_clip=False,
        bbox=dict(
            ec =color, 
            boxstyle='round,pad=0.2', 
            fc = "w")
            )
    #add the other text
    ax_matrix.text(
        2.3,
        -4.5,
        f'since 2004' ,
        color = "#101628", 
        size = 10, 
        )

    #add some styling
    ax_matrix.set_frame_on(False)
    ax_matrix.tick_params(
        axis='both', 
        which='major', 
        length=0,
        labelleft = False,
        labelbottom = False
         )