import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from matplotlib.lines import Line2D
from matplotlib.offsetbox import OffsetImage, AnnotationBbox

color_dict = { 2022: "#A54836", 2004: "#5375D4", }

xy_ticklabel_color, legend_color, grid_color, datalabels_color ='#757C85',"#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":3, "Norway":1, 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.year.map(color_dict)
df

fig, ax= plt.subplots(figsize=(5,5), facecolor = "#FFFFFF",sharex=True, sharey=True, subplot_kw=dict(polar=True))

ax.set_theta_zero_location('N')
ax.spines["polar"].set_color("none")
ax.set_rlim(0,20)

for i, (year, group) in enumerate(df.groupby("year")):
    sites = group['sites'].tolist()
    sites_radar = sites + [sites[0]]

    angles = np.linspace(0, 2*np.pi, len(countries), endpoint=False)
    angles_radar = np.concatenate((angles,[angles[0]])) 

    #print(sites_radar)
    radar_plot = ax.plot(
        angles_radar, 
        sites_radar, 
        'o', 
        ms= 8,
        mec="w",
        linewidth=2, 
        color = group['color'].iloc[0],
        zorder = 2, 
        clip_on=False
        )

    #add the arcs.    
    sites_start = sites[1:] + [sites[0]]  #shift the first to last to create the connection points
    angle_start = np.concatenate((angles[1:], [angles[0]])) 

    for i in range(3):
        ax.annotate(
            "", 
            xy=(angle_start[i], sites_start[i]), 
            xytext=(angles[i], sites[i]), 
            zorder = 1,
            arrowprops=dict(
                arrowstyle='-', 
                connectionstyle='arc3,rad=0.2', 
                color = group['color'].iloc[0],  
                linewidth=1, 
                linestyle='-', 
                antialiased=True
                )
            )

# Axis labels
ax.set_xticks(angles_radar[:-1])
ax.set_xticklabels([])

# Hide default radial ticks
ax.set_yticks([])
ax.set_ylim(0, 15)

# Tick and label settings
tick_radii = np.arange(1, 16, 1)     # ticks every 1 unit
label_radii = np.arange(5, 20, 5)    # label every 5 units
tick_length = 0.8                    # tick mark length
label_offset = 1.6                   # label distance from tick (along fixed direction)

for angle in angles_radar[:-1]:
    for r in tick_radii:
        # Coordinates on the spoke
        x = r * np.cos(angle)
        y = r * np.sin(angle)

        # Perpendicular direction for tick
        tick_perp_angle = angle + np.pi / 2
        dx = (tick_length / 2) * np.cos(tick_perp_angle)
        dy = (tick_length / 2) * np.sin(tick_perp_angle)

        # Tick line start and end
        x0, y0 = x - dx, y - dy
        x1, y1 = x + dx, y + dy

        # Convert to polar for tick plotting
        r0 = np.hypot(x0, y0)
        theta0 = np.arctan2(y0, x0)
        r1 = np.hypot(x1, y1)
        theta1 = np.arctan2(y1, x1)

        ax.plot(
            [theta0, theta1],
             [r0, r1], 
             color=xy_ticklabel_color, 
             linewidth=.5
             )

        # Add label (only for selected ticks)
        if r in label_radii:
            label_angle_offset = angle + np.pi / 2 if np.cos(angle) >= 0 else angle - np.pi / 2


            # Apply offset in angular direction (fixed)
            lx = x - label_offset * np.cos(label_angle_offset)
            ly = y - label_offset * np.sin(label_angle_offset)

            r_label = np.hypot(lx, ly)
            theta_label = np.arctan2(ly, lx)

            ax.text(
                theta_label, r_label, str(r),
                ha='center',
                va='center',
                fontsize=9,
                color=xy_ticklabel_color,
                rotation=0
            )

fig

for ang, im in zip(angles,  img):
    imagebox = OffsetImage(im, zoom=.04)
    ax.add_artist(AnnotationBbox(
        imagebox, 
        (ang, df.sites.max()+2), 
        frameon = False, 
        annotation_clip=False
        )
    )
fig

labels = df.year.unique()

lines = [Line2D([0], [0], color=c,  marker='o',linestyle='', markersize=8,) for c in df.color.unique()]
fig.legend( 
    lines,
    labels, 
    labelcolor=legend_color,
    bbox_to_anchor=(0.5, 0), 
    loc="lower center",
    ncols = 2,
    frameon=False, 
    fontsize= 12
    )

ax.text(
    0,
    0, 
    "▼", 
    color = xy_ticklabel_color, 
    ha="center", 
    va="center"
    )

fig