#tutorial 
import matplotlib.pyplot as plt
from matplotlib.patches import Wedge
from matplotlib.lines import Line2D
import matplotlib.patches as mpatches
import pandas as pd
import numpy as np

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

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

data = {
    "year": [2004, 2022, 2004, 2022, 2004, 2022],
    "countries" : [ "Denmark", "Denmark", "Norway", "Norway","Sweden", "Sweden"],
    "sites": [4,10,5,8,13,15]
}
df= pd.DataFrame(data)


df['sub_total'] = df.groupby('countries')['sites'].transform('sum')
df['pct_change'] = df.groupby('countries', sort=False)['sites'].apply(
     lambda x: x.pct_change()).to_numpy()*-1
df = df.sort_values(['year','sites' ], ascending=[True,False] )
#map the colors of a dict to a dataframe
df['color']= df.countries.map(color_dict)
df

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

ax.set_xlim([-2, site_max + 2])
ax.set_ylim([-8, 8])

for i, row in enumerate(df.itertuples()):
    center = (row.sites/2, 0) 
    radie = (row.sites/2)
    
    dir = direction[i]
    #print(center, radius)
    #add the arcs
    arc = Wedge(
        center, 
        radie * dir, 
        theta1, 
        theta2, 
        color=row.color,
        )
    ax.add_patch(arc)

    ax.text(
        center[0] + radie -.8,
        .55 * dir, 
        row.sites, 
        ha = "left",
        va = "center", 
        color = "w"
    )

ax.axhline(
    y=0,
    xmin=-2,
    xmax=site_max, 
    color =grid_color,
    lw=0.5,
    snap=False
    )

year = df.year.unique().tolist()
y_coords = [1, -2]
for i, y in enumerate(y_coords):
    ax.text(
        -4, 
        y, 
        year[i], 
        color=xy_ticklabel_color, 
        size=12 
    )

fig

# Legend axis at the bottom center
legend_ax = fig.add_axes([0.2, -0.15, 0.6, 0.4])  
legend_ax.set_aspect('equal')  # Ensure wedges look circular
legend_ax.axis("off")

legend_ax.set_xlim(0, 2) 

sort_order_dict = {"Denmark": 3, "Sweden": 2, "Norway": 1}
legend_sort = df.sort_values(by="countries", key=lambda col: col.map(sort_order_dict))

# Semicircle + label data
wedge_centers = [(0.25, 0.2), (0.75, 0.2), (1.25, 0.2)]

# Create wedges and labels
for center, color, label in zip(wedge_centers, legend_sort.color.unique(), legend_sort.countries.unique()):
    wedge = mpatches.Wedge(
        center=center,
        r=0.08,           # radius
        theta1=theta1,
        theta2=theta2,
        facecolor=color
    )
    legend_ax.add_patch(wedge)

    # Add label to the right of each wedge
    legend_ax.text(
        center[0] + 0.12, center[1] -0.01,  # slightly to the right
        label,
        va='bottom', 
        ha='left',
        fontsize=9, 
        color=xy_ticklabel_color
    )

 
ax.set_frame_on(False)
ax.tick_params(
    axis='both', 
    which='both',
    length = 0, 
    labelleft =False, 
    labelbottom = False
)

fig