#tutorial 
import matplotlib.pyplot as plt
from matplotlib.patches import Wedge
from itertools import cycle
import pandas as pd
import numpy as np

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

code_dict = {"Norway": "NO", "Denmark": "DK", "Sweden": "SE" }

xy_ticklabel_color, xlabel_color,  grid_color, datalabels_color ='#757C85',"#101628", "#C8C9C9", "#FFFFFF"

data = {
    "year": [2004, 2022, 2004, 2022, 2004, 2022],
    "countries" : ["Norway", "Norway","Sweden", "Sweden",  "Denmark", "Denmark"],
    "sites": [5,8,13,15,4,10,]
}
df= pd.DataFrame(data)
df = df.sort_values(['year' ], ascending=False ).reset_index(drop=True)
df['ctry_code'] = df.countries.map(code_dict)
#Add the color based on the color dictionary
df['color'] = df.set_index(['year', 'countries']).index.map(color_dict.get)

df

offset = 10

# Get site counts for the latest year
sites = df.loc[df.year == df.year.max(), 'sites'].to_numpy()

# Compute centers directly using vectorized operations
radii = sites / 2
centers = np.cumsum(sites) - radii

#add the offset incrementally
adjusted_centers = centers + np.arange(len(centers)) * (2 * offset) 

# Create a flat list with adjusted centers repeated twice
y_coord = np.tile(adjusted_centers, 2).tolist()
print(y_coord)

fig, ax = plt.subplots( figsize=(10,5), facecolor = "#FFFFFF")
ax.set_aspect('equal')

for i, row in enumerate(df.itertuples()):
    dir_ = direction[i]
    y = y_coord[i]

    patches = Wedge(
        (0, y), 
        row.sites , 
        90 * dir_, 
        -90 * dir_, 
        color=row.color, 
        zorder = 1,
        )
    ax.add_patch(patches)

    ax.text(
        -25 * dir_, 
        y, 
        row.sites, 
        fontsize=10, 
        color = xy_ticklabel_color,
        ) 
    ax.text(
        -2, 
        y, 
        row.ctry_code, 
        fontsize=10, 
        color = "w",
        va="center",
        weight = "bold"
        ) 
    
    if i in(0,3):
        ax.text(
            25 * dir_, 
            81, 
            row.year, 
            size= 12, 
            ha="center", 
            color= xlabel_color
    )

ax.set_xlim([-30, 30])
ax.set_ylim([-10, 80])

ax.axvline(
    x=0,
    ymin=-4,
    ymax=17, 
    color =grid_color,
    lw=0.5,
    snap=False,
    zorder = 0,
)

ax.set_frame_on(False)
ax.tick_params(labelleft = False, labelbottom = False, length =0)

fig