import matplotlib.pyplot as plt
from matplotlib.lines import Line2D
import numpy as np
import pandas as pd


color_dict = {(2022,"Norway"): "#9194A3", (2004,"Norway"): "#2B314D",
              (2022,"Denmark"): "#E2AFA5", (2004,"Denmark"): "#A54836",
              (2022,"Sweden"): "#C4D6F8", (2004,"Sweden"): "#5375D4",
              }

xy_ticklabel_color, legend_color,  datalabels_color ="#101628", "#757C85", "#FFFFFF"

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)
df['sub_total'] = df.groupby('countries')['sites'].transform('sum')
df['total'] = 20  #the circle has been divided by 20 parts
df['pct_group'] = df['sites'] / df.total

#custom sort
sort_order_dict = {"Denmark":2, "Sweden":3, "Norway":1, 2004:5, 2022:4}
df = df.sort_values(by=['countries','year'], key=lambda x: x.map(sort_order_dict))

#Add the color based on the color dictionary
df['color'] = df.set_index(['year', 'countries']).index.map(color_dict.get)

df

fig, ax = plt.subplots(figsize=(10,5), facecolor = "#FFFFFF", subplot_kw=dict(polar=True) )

ax.set_theta_zero_location('N')
ax.set_theta_direction(-1)
ax.set_axis_off() 

for i, (country, group) in enumerate(df.groupby("countries")):
    for row in group.itertuples():
        angle_range = np.linspace(0, row.pct_group * 2 * np.pi)
        r = np.full(len(angle_range), i + 1)  # identical radius values to draw an arc
        line = ax.plot(
            angle_range,
            r,
            linewidth=15,
            solid_capstyle="round",
            color=row.color
            )
        
        ax.annotate(
            row.sites, 
            xy= ( angle_range[-1],r[-1]), 
            color=datalabels_color,
            ha="center" ,
            va="center"
            )

#x position in cartesian coordinates
x =[.5]*3
for x,y, country in zip(x,range(1,4), df.countries.unique()):
    r = np.sqrt(x**2 + y**2)
    theta = np.arctan2(y, x)
    ax.text(
        theta+np.deg2rad(-90),
        r, 
        country, 
        ha= "right", 
        va="center" 
        )
     
ax.set_rmax(4)

fig

text_legends = ["Before", "After"]
years = df.year.unique()

lines = [Line2D([0], [0], color=c, linestyle='-', ) for c in reversed(df.color.unique())]
labels  = [f'{text_legend} 2004' for year, text_legend in zip(years , text_legends)]

for year in years:
    plt.figlegend( 
        lines,
        labels,   
        labelcolor=xy_ticklabel_color,
        bbox_to_anchor=(0.5, 0), 
        loc="lower center",
        ncols = 2,
        frameon=False, 
        fontsize= 10
        )

fig