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

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

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)
df = df.sort_values(
    [
        "year",
        "sites",
    ],
    ascending=True,
).reset_index(drop=True)

df["sub_total"] = df.groupby("year")["sites"].transform("sum")
df["pct_group"] = df["sites"] / df.sub_total

df["ctry_code"] = df.countries.map(code_dict)
# map the colors of a dict to a dataframe
df["color"] = df.countries.map(color_dict)

df

fig, axes = plt.subplots(ncols=df.year.nunique(), figsize=(8, 8), sharey=True, facecolor="#FFFFFF", subplot_kw=dict(polar=True),)
fig.tight_layout(h_pad=-5)

heights = [0.6, 1]  # Define the height for each bar
directions = [1, -1] #Define the direction of the chart
offsets = [-1, -2]  #Define how much to offset
for i, ((year,  group), ax) in enumerate(zip(df.groupby("year"), axes.ravel())):

    degrees = group["pct_group"] * 180 #half pie
    bottom = np.zeros(1)
    for deg, color in zip(degrees,  group.color):
        #print(height)
        ax.barh(
            range(1), 
            np.radians(deg), 
            height=heights[i], 
            left=bottom, 
            color=color
            )
        bottom += np.radians(deg)
    
    #add the bar data labels
    for bar, row in zip(ax.patches, group.itertuples()):
        x = bar.get_x() + bar.get_width() / 2
        y = bar.get_height() / 2 + bar.get_y()
        ax.text(
            x,
            y,
            f"{row.ctry_code}\n{row.sites}",
            ha = "center",
            va = "center",
            color = datalabels_color,
            size = 10,
        )

    #add the subttotals
    ax.text(
        0.5,
        0.5,
        group.sub_total.iloc[0],
        size=22,
        ha="center",
        va="center",
        transform=ax.transAxes,
        color=xy_ticklabel_color
    )

    
    ax.set_title(
        year, 
        color = xy_ticklabel_color
        )

    ax.set_theta_zero_location("N")
    ax.set_theta_direction(directions[i])
    ax.set_rorigin(offsets[i])

    shift_axes = 0.2 if directions[i] == 1 else -0.2
    
    box = ax.get_position()
    box.x0 = box.x0 + shift_axes  # x0 first coordinate of the box
    box.x1 = box.x1 + shift_axes  # x1 last coordinate of the box
    ax.set_position(box)
    ax.set_axis_off()


line = plt.Line2D((0.5, 0.5), (0.2, 0.8), color=grid_color, linewidth=1)
fig.add_artist(line)