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

color_dict = {"Norway": "#2B314D", "Denmark": "#A54836", "Sweden": "#5375D4"}

xy_ticklabel_color, grand_totals_color, pie_color, datalabels_color = (
    "#757C85",
    "#101628",
    "#ECEFEF",
    "#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("year")["sites"].transform("sum")

sort_order_dict = {"Denmark": 2, "Sweden": 1, "Norway": 3, 2004: 4, 2022: 5}
df = df.sort_values(
    by=[
        "year",
        "countries",
    ],
    key=lambda x: x.map(sort_order_dict),
)
# map the colors of a dict to a dataframe
df["color"] = df.countries.map(color_dict)

df

unique_countries = df.countries.unique()
unique_years = df.year.unique()

grey_circles = [df.sites.max()] * len(df)
print(grey_circles)

fig, axes = plt.subplots(nrows=len(unique_years), ncols=len(unique_countries), figsize=(10, 5), facecolor="#FFFFFF", sharex=True, sharey=True, subplot_kw=dict(polar=True))

for row, grey_circle, ax in zip(df.itertuples(), grey_circles, axes.ravel()):
    ax.set_rmax(7)
    ax.set_theta_zero_location("N")
    ax.set_theta_direction(-1)


    angle_range = np.linspace(0, row.sites * 25)
    angle_grey_cicle = np.linspace(0, grey_circle * 25)
    r = np.full(len(angle_range), 6)  # identical radius values to draw an arc

    ax.plot(
        [np.deg2rad(a) for a in angle_grey_cicle],
        r,
        lw=10,
        color=pie_color,
    )
    ax.plot(
        [np.deg2rad(a) for a in angle_range],
        r,
        linewidth=10,
        solid_capstyle="round",
        color=row.color,
    )

    ax.annotate(
        text=row.sites,
        xy=(np.deg2rad(0), r[-1]),
        xytext=(0, -60),
        size=24,
        weight="bold",
        color=grand_totals_color,
        textcoords="offset points",
        ha="center",
        va="center",
    )

    ax.grid(False)
    ax.set_frame_on(False)
    ax.tick_params(labelleft = False, labelbottom = False)

for ax, year in zip(axes[:, 0], unique_years):
    ax.set_ylabel(
        year,
        rotation=0,
        size=14,
        color=xy_ticklabel_color,
        labelpad = 53
    )


for ax, col in zip(axes[0], unique_countries):
    ax.set_title(
        col,
        x=0.5,
        y=1.2,
        color=xy_ticklabel_color,
    )
fig