import matplotlib.pyplot as plt
import math
import numpy as np
import pandas as pd


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

xy_ticklabel_color, grand_totals_color, grid_color, datalabels_color ='#757C85',"#101628", "#ECEFF1", "#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)
#custom sort
sort_order_dict = {"Denmark": 2, "Sweden": 3, "Norway": 1, 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), figsize=(10, 2))

for row, shift_axe, ax in zip(df.itertuples(), shift_axes, axes.ravel()):
    ax.set(xlim = (-2,2), ylim  = (-0,1))

    x = y = np.arange(1)
    site = row.sites    
    #add the triangles
    ax.scatter(
        x,
        y, 
        s=site*400, 
        marker=10, 
        color = row.color,  
        zorder=1
        )

    #add vertical line
    ax.axvline(
        x=0, 
        ymin = 0, 
        ymax =1, 
        color=grid_color, 
        linestyle= "-",
        clip_on = False, 
        zorder= 0
        ) 

    #add site labels
    ax.annotate(
        site, 
        (x, y+ math.sqrt( site ) /10  ), 
        size= 10, 
        ha= "center", 
        va="top", 
        color= "w"
        )

    #shift axes
    box = ax.get_position()
    box.x0 = box.x0 + shift_axe #x0 first coordinate of the box
    box.x1 = box.x1 + shift_axe #x1 last coordinate of the box
    ax.set_position(box)
    ax.patch.set_alpha(0.1) #transparent axis

    ax.set_frame_on(False)
    ax.tick_params(
        length=0, 
        labelleft = False, 
        labelbottom = False
        ) 
    ax.set_title(
        row.year, 
        size= 9, 
        color = xy_ticklabel_color
        )

x = [0.2,0.3,0.43]
for x,country in zip(x,df.countries.unique()):
    ax.text( 
        x,
        -0.1, 
        country, 
        size = 10,
        color = xy_ticklabel_color, 
        transform = fig.transFigure
        )