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


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

xy_ticklabel_color, grid_color, datalabels_color = "#101628", "#C8C9C9", "#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)
#custom sort
sort_order_dict = {"Denmark":2, "Sweden":1, "Norway":3, 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['diff'] = df.groupby(['countries'])['sites'].diff()*-1
df['diff'] = df['diff'].fillna(0)
df

fig, ax = plt.subplots(figsize=(8, 8), facecolor="#FFFFFF")

countries = df.countries.unique()
widths = [0, 0.25] 
ax.set(ylim = (0,3))

for i, (year,  group) in enumerate(df.groupby("year", sort = False)):
    new_sites = df[df['year'] == year][['diff', 'sites']].values.flatten().tolist()
    new_sites.append(sum(new_sites)) 

    color = ["#FFFFFF"] * (len(group.color) * 2) 
    color[1::2] = group['color']
    color.append("#FFFFFF")

    wedges, texts = ax.pie(
        new_sites,
        radius=1 - widths[i],
        startangle=-90,
        wedgeprops=dict(width=0.2),
        colors=color,
    )

    labels = [""] * (len(group.color) * 2) 
    labels[1::2] = group['sites']
    labels.append("")

    shift = 5
    for lbl, wedge in zip(labels, wedges):
        #shift from the start angle
        angle_lbls = wedge.theta2 - shift
        # center of wedge
        r = wedge.r - wedge.width / 2
        # convert polar to cartesian
        x_lbl = r * np.cos(np.deg2rad(angle_lbls))
        y_lbl = r * np.sin(np.deg2rad(angle_lbls))

        ax.annotate(
            lbl,
            xy=(x_lbl, y_lbl),
            size=12,
            color=datalabels_color,
            ha="center",
            va="center",
            weight="bold",
        )
    

    # Add the Year legend
    for year, w in zip([year], wedges):
        angle = w.theta2
        r = w.r - w.width / 2
        x = r * np.cos(np.deg2rad(angle))  
        y = r * np.sin(np.deg2rad(angle))  
        ax.annotate(f"{year}  ", xy=(0, r), size = 14, ha="right", va= "center", color=xy_ticklabel_color)

# Add the country labels
label_spacing = 0.25
country_labels = {2: "Sweden", 4: "Denmark", 6: "Norway"}
country_colors = df[df.year == 2022].color.unique()

for key, color in zip(country_labels, country_colors):
    w = wedges[key - 1]
    angle = w.theta2
    r = w.r + w.width + label_spacing
    x = r * np.cos(np.deg2rad(angle))
    y = r * np.sin(np.deg2rad(angle))

    ang = np.deg2rad((w.theta1 + w.theta2) / 2)
    y_line = np.sin(ang)
    x_line = np.cos(ang)

    #add country labels
    ax.annotate(
        country_labels[key] + " ",
        xy=(x, y-.1),
        size=14,
        color=color,
        ha="left",
        va="center",
    )

    #lines at each country
    ax.annotate(
        "",
        xy=(0.4 * x, 0.4 * y),
        xytext=(0.85 * x, 0.85 * y),
        arrowprops=dict(arrowstyle="-", color=grid_color, lw=0.5),
        va="center",
    )


ax.axvline(0, 0.1, 0.7, color=grid_color, lw=0.5)
fig