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

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

xy_ticklabel_color, title_color, grid_color, datalabels_color ='#101628',"#101628", "#C8C9C9", "#101628"

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)

#custom sort
sort_order_dict = {"Denmark":2, "Sweden":3, "Norway":1, }
df = df.sort_values(by=['countries',], 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['ctry_code'] = df.countries.map(code_dict)
df

offset_labels=[0.4,0.3,0.2]
directions = [1,-1]

fig, axes = plt.subplots(ncols=2, figsize=(10,5), facecolor = "#FFFFFF", subplot_kw=dict(polar=True) )
fig.tight_layout(h_pad=-40)

for ax, (year, group), direction, offset in zip(axes.ravel(),df.groupby("year"), directions, offset_labels):

   
    ax.set_theta_zero_location('N')
    ax.set_theta_direction(direction)

    for i,row in enumerate(group.itertuples()):
        angle_range = np.linspace(0, row.sites*7)
        theta =[np.deg2rad(a) for a in angle_range]
        r = np.full(len(angle_range), i + 1)  
        ax.plot(
            theta,
            r,
            linewidth=15,
            solid_capstyle="round",
            color=row.color
            )

        ax.annotate(
            row.sites, 
            xy= ( theta[-1]+offset,r[-1]), 
            color=datalabels_color,
            ha="center" ,
            va="center"
            )

    ax.set_title(
        year, 
        size= 16, 
        weight = "bold", 
        color = title_color
        )
    ax.set_rmax(4)
    ax.grid(False)
    ax.axis('off')

    #shift axes
    shift_axes = 0.18 if direction == 1 else -0.18
    box = ax.get_position()
    box.x0 = box.x0 + shift_axes
    box.x1 = box.x1 + shift_axes
    ax.set_position(box)

#add country legend
x= [1.5]*3
for x,y, code in zip(x,range(1,4), df.ctry_code):
    r = np.sqrt(x**2 + y**2)
    theta = np.arctan2(y, x)
    ax.text(
        theta+np.deg2rad(-90),
        r, 
        code , 
        va= "center")

line = plt.Line2D((.5,.5),(.85,1.), color=grid_color, linewidth=1)
fig.add_artist(line) 

fig