import matplotlib.pyplot as plt
from matplotlib.offsetbox import OffsetImage, AnnotationBbox
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",
}

xlabel_color, title_color, datalabels_color, line_color  = "#101628", "#969CA3", "#2B314D", "#DFE4E5"

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["year_lbl"] = "'" + df["year"].astype(str).str[-2:].astype(str)
sort_order_dict = {"Denmark": 2, "Sweden": 3, "Norway": 1, 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["diff"] = df.groupby(["countries"])["sites"].diff().abs()
df["diff"] = np.where(df["diff"].isnull(), df.sites, df["diff"]).astype(int)
df

fig, axes = plt.subplots(nrows=df.countries.nunique(), ncols=1, figsize=(5, 7), facecolor="#FFFFFF", subplot_kw=dict(polar=True))
fig.tight_layout(pad=-4.8)

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

    direction = plot_direction[j]
    ax.set_theta_zero_location("N")
    ax.set_theta_direction(direction)
    ax.set_axis_off()

    # plot the radial lines
    sites_max = group["sites"].max() 
    ax.set_ylim(0, sites_max + 1)

    # number of total angles 
    total_angles = sites_max * 2
    lines_to_draw = sites_max  # Only draw the first 10

    # generate the angles (in radians) for the full circle
    angles = np.linspace(0, 2 * np.pi, total_angles, endpoint=False)

    #create the bubble colors 
    bubble_color = [group['color'].iloc[0]]*group['diff'].iloc[0] + [group['color'].iloc[1]]*group['diff'].iloc[1] 
    # Draw only the first 10 radial lines
    for i, theta in enumerate(angles[:lines_to_draw]):   
        #print(group['sites'].iloc[0])     
        ax.plot(
            [theta, theta], 
            [0, sites_max], 
            color= line_color, 
            linewidth=1,
            zorder = 0
            )

        #add the bubbles
        ax.scatter(
            theta, 
            sites_max, 
            s = 150,
            color =bubble_color[i], 
            ec="w", 
            lw = 3,
            zorder=1
        )
    # add flag
    image_box = OffsetImage(img[j], zoom=0.04)  # container for the image
    ab = AnnotationBbox(image_box, (0, 0), frameon=False)
    ax.add_artist(ab)

    #add the  white border to the flag
    ax.scatter(
        0, 
        0, 
        s = 650,
        color ="w", 
        zorder=1
    )

    ha_ctry = 'left' if direction == 1 else 'right'    
    #add country names 
    ax.text(
        0.5 + .1 * direction,
        0.5,
        f"{group['countries'].iloc[0]}",
        color=title_color,
        transform=ax.transAxes,
        ha=ha_ctry,
        va="center",
    )

    #add subtotals
    ax.text(
            0.5 + .35 * direction,
            0.5,
            f"{sites_max}",
            size = 10,
            color=group['color'].iloc[0],
            weight= "bold",
            transform=ax.transAxes,
            ha=ha_ctry,
            va="center",
        )

    #add year labels
    ha_lbl = 'right' if direction == 1 else 'left'  #change the direction of the labels depending on the chart orientation
    for k, row in enumerate(group.itertuples()):
        ax.text(
            angles[row.sites - 1],
            sites_max + 1,
            row.year_lbl,
            ha=ha_lbl, 
            va='top', 
            fontsize=8,
            color=row.color
            )