import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from matplotlib.offsetbox import OffsetImage, AnnotationBbox

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

xy_ticklabel_color, xlabel_color, grand_totals_color, grid_color, datalabels_color ='#101628',"#101628","#101628", "#C8C9C9", "#2B314D"

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

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

fig, axes = plt.subplots(ncols =3, figsize=(10,5), sharex=True, sharey=True, facecolor = "#FFFFFF", subplot_kw=dict(polar=True) )
fig.tight_layout(pad=3.0)

for (country, group), im, ax in zip(df.groupby("countries", sort = False), img, axes.ravel()):

    ax.set_theta_zero_location("N")
    ax.set_theta_direction(-1)
    ax.set_rorigin(8)

    #add images
    image_box =  OffsetImage(im, zoom = 0.05) #container for the image
    ab = AnnotationBbox(image_box, (0.5, 0.5), frameon=False, xycoords='axes fraction', boxcoords="axes fraction")
    ax.add_artist(ab)

    #add bars
    diff_subtotal = (group['sub_total']-15).abs()
    
    diffs =  np.append(group['diff'].values, diff_subtotal.unique())    
    bar_color = np.append(group['color'].values,"#ECEFEF" )

    bottom =  np.zeros(1)
 
    for i, diff in enumerate(diffs):
        angles = np.radians(diff/15 *360)
        bar = ax.barh(
            range(1), 
            angles,  
            left=bottom, 
            color=bar_color[i]
            )
        bottom += angles



        ax.set_title(
            group['countries'].unique().item(), 
            color = xy_ticklabel_color, 
            x=0.5, 
            y =1.3, 
            size= 14, 
            weight = "bold"
            )

    for bar,rows in zip(ax.patches,group.itertuples()):
        x= bar.get_x() + bar.get_width()
        y = bar.get_height()/2 + bar.get_y() 
        print()
        #add data labels
        ax.text( 
            x, 
            y -2, 
            rows.sites, 
            ha='center', 
            color=xy_ticklabel_color,  
            size=10
            )    
        ax.axvline( 
            x, 
            y, 
            y - 1, 
            color=rows.color, 
            lw=1,
            clip_on= False 
            )
        ax.set_axis_off()