import matplotlib.pyplot as plt
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 ='#101628',"#101628", "#C8C9C9", "#101628"

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['ctry_code'] = df.countries.astype(str).str[:2].astype(str).str.upper()

#custom sort
sort_order_dict = {"Denmark":2, "Sweden":1, "Norway":3, 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.year.unique()), nrows = len(df.countries.unique()),sharex= False, figsize=(15, 15), subplot_kw=dict(polar=True))
fig.tight_layout(h_pad=-20)

shift_axes= [+0.1, -0.1]*3
for ax, row in zip(axes.ravel(), df.itertuples()):
    for site in range(row.sites):
        angles = np.linspace(0,np.pi, 100)
        radius = np.ones(100) * site

        ax.plot(
            angles, 
            radius, 
            color=row.color, 
            linestyle='-',
            lw=1
            )

        ax.set_xlabel(
            row.sites, 
            size=20, 
            color=xy_ticklabel_color, 
            weight="bold", 
            labelpad = -80
            )

        ax.hlines(
            y=0.25, 
            xmin=0.0, 
            xmax=1.2,
            colors='lightgrey', 
            linewidth=1, 
            linestyles='-',
            transform=ax.transAxes, 
            clip_on = False
        )
        
        ax.scatter(
            0,
            0, 
            s = 30,
            color = row.color,
            clip_on = False,
            zorder = 10,
            )

        #styling
        ax.set_rgrids(range(0,15))
        ax.grid(False)
        ax.set_frame_on(False)
        ax.tick_params(labelleft = False, labelbottom = False)
        ax.set_thetamax(180)


  

for ax, code in zip(axes[:,0], df.ctry_code.unique()):
    ax.set_ylabel(
        code, 
        rotation=0,
        size= 20, 
        color=xy_ticklabel_color, 
        labelpad  = 33
        )
    ax.yaxis.set_label_coords(-0.2, 0.22)

for ax, year in zip(axes[0], df.year.unique()):
    ax.set_title(
        year, 
        x=0.5, 
        y=-1.5, 
        size = 20, 
        color=xy_ticklabel_color
        )
fig