import geopandas as gpd
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
from matplotlib.patches import Arc


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

xy_ticklabel_color, map_color, grid_color, datalabels_color, h_line ='#757C85',"#D3D3D3", "#C8C9C9", "#FFFFFF", "#939AA1"

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['diff']=df.groupby('countries')['sites'].diff().fillna(df.sites).astype(int)

sort_order_dict = {"Denmark": 2, "Sweden": 3, "Norway": 1, 2004: 5, 2022: 4}
df = df.sort_values(by=["countries","year"], key=lambda x: x.map(sort_order_dict))

df['sub_total'] = df.groupby('countries')['diff'].transform('sum')

#map the colors of a dict to a dataframe
df['color']= df.year.map(color_dict)
df['pct_subtotal'] = df['diff'] / df['sub_total']

df

# Get the NUTS2 GeoJSON here: https://github.com/eurostat/Nuts2json
url = "https://raw.githubusercontent.com/eurostat/Nuts2json/master/pub/v2/2021/3035/20M/0.json"
map_df = gpd.read_file(url,  layer='nutsrg')
map_df = map_df[map_df.id.isin(['NO','SE', 'DK']) ]

map_df

fig = plt.figure(figsize=(15,10))
ax_map = fig.add_axes([0, 0, 1, 1])
map_df.plot(color=map_color,ax=ax_map)
ax_map.set_axis_off()

for i, (country, group) in enumerate(df.groupby("countries", sort = False)):
    axes_height = 0.2
    axes_width = 0.2
    #add the axes at y and x
    ax_pie = fig.add_axes([x[i], y[i] , axes_width, axes_height]) 

    #add the pies
    patches, texts, autotexts  =  ax_pie.pie(
        group['pct_subtotal'], 
        autopct='%1.0f%%',
        pctdistance=0.75,
        wedgeprops=dict(width=0.5),    
        startangle=90,
        colors= group.color
        )
    
    for autotext in autotexts:
        autotext.set_fontsize(12)           
        autotext.set_color('w')     


    #add the sub totals    
    ax_pie.text(
        np.pi/2-1.6,
        0, 
        group.sub_total.iloc[0],
        va="center", 
        ha="center", 
        size= 26
        )
    #add country labels
    ax_pie.text(
        np.pi/2-1.6,
        -1.3, 
        group.countries.iloc[0],
        va="center", 
        ha="center", 
        size= 12 
        )
    #add horizontal line
    ax_pie.axhline( 
        y = 0, 
        xmin = xmin_line[i], 
        xmax = xmax_line[i],
        lw=1, 
        clip_on = False, 
        color=h_line
        )

    # hlines uses axes coordinates, convert to data coordinates
    x0, x1 = ax_pie.get_xlim()
    x_end = x0 + xmax_line[i] * (x1 - x0)

    # add a circle at the end using scatter
    ax_pie.scatter(
        x_end,            
        y_line + .4,           
        s=20,             
        color=h_line,  
        fc = map_color,   
        zorder=3,         
        clip_on=False     
    )

    
    # Add an arch on top of the pie
    arc = Arc(
        xy=(0, 0),            
        width=2.5, 
        height=2.5,
        angle=0,                # rotation of the arc (0 means aligned with x-axis)
        theta1=arc_theta1[i], 
        theta2=arc_theta2[i],  # start and end angles in degrees (relative to 0 = east)
        edgecolor=h_line,
        lw=1, 
        clip_on = True
    )
    ax_pie.add_patch(arc)
 
fig

# Legend axis setup
legend_ax = fig.add_axes([.65, -0.1, 0.08, 0.08])
legend_ax.axis("off")

# Semicircle ring (wedge with inner radius)
semicircle_data = [
    {
        "center": (0.4, 0.3), "r": 0.15, "width": 0.07,
        "theta1": 90, "theta2": 270,  # Left-facing
        "facecolor": color_dict[df.year.max()], 
    },
    {
        "center": (0.45, 0.3), "r": 0.15, "width": 0.07,
        "theta1": -90, "theta2": 90,  # Right-facing
        "facecolor": color_dict[df.year.min()], 
    }
]


for wedge_kwargs in semicircle_data :
    wedge = mpatches.Wedge(**wedge_kwargs)
    legend_ax.add_patch(wedge)

# Text data
text_data = [
    {"x": -0.4, "y": 0.3, "s":"After 2004", "size": 8, "va": "center", "color": h_line},
    {"x": .7, "y": 0.3, "s":"Before 2004", "size": 8, "va": "center", "color": h_line},
]

# Add text
for txt in text_data:
    legend_ax.text(**txt)

fig