import matplotlib.pyplot as plt
from matplotlib.offsetbox import OffsetImage, AnnotationBbox
import matplotlib.transforms as transforms
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, xlabel_color,  grid_color,  ='#9BA0A6',"#9BA0A6", "#E9ECED", 


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 = df.sort_values([ 'year'], ascending=True ).reset_index(drop=True)
df['diff'] = df.groupby(['countries'])['sites'].diff()
df['diff'] = df['diff'].fillna(df.sites)
df['sub_total'] = df.groupby('countries')['diff'].transform('sum')
df['pct_group'] = (df['diff'] / df.sub_total)*100
df['pct_group'] = df['pct_group'].astype(int)
#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))

#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=(5,5), facecolor = "#FFFFFF")

ax.xaxis.set_ticks(np.arange(0, 20, 5), )
ax.yaxis.set_ticks(np.arange(0, 20, 5), )


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

    x = group.sites.iloc[0]
    y = group.sites.iloc[1]
    colors = group.color.tolist()
    #create a new axes for each pie with the coordinates of ax
    ax_pie = ax.inset_axes([x, y, 0.2, 0.2], transform=ax.transData) 

    wedges, texts = ax_pie.pie(
        group["diff"],
        radius = 15,
        wedgeprops = dict(
            width=8),  
        labels = [f"{p:.0f}%" for p in group.pct_group], 
        labeldistance = 1.2,
        textprops = {
            'fontsize':  8, 
            'color': colors[0],
            },
         counterclock = False,
         startangle = 90, 
         colors = group.color
         )

    #add the images
    offset = .1
    image_box =  OffsetImage(im, zoom = 0.03) #container for the image
    ab = AnnotationBbox(
        image_box, 
        (x + offset, y + offset), 
        frameon = False,
        annotation_clip = False
        )
    ax.add_artist(ab)

ax.tick_params(
    axis='both', 
    which='major',
    length=0, 
    labelsize=12,
    colors =xy_ticklabel_color
    )

ax.grid(color = grid_color)

label_params = dict(
    size=12, 
    color = xlabel_color, 
    weight= "bold"
    )

ax.set_xlabel( df.year.min(), **label_params)
ax.set_ylabel(df.year.max(), **label_params)

for axis in ['top', 'bottom', 'left', 'right']:
    ax.spines[axis].set_color(grid_color) 
    ax.spines[axis].set_zorder(0)

fig