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

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

xy_ticklabel_color, grand_totals_color, spines_color, datalabels_color ='#757C85',"#101628", "#828696", "#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)

df['sub_total'] = df.groupby('year')['sites'].transform('sum')
df['diff'] = df.groupby(['countries'])['sites'].diff()
#custom sort
df = df.sort_values(by=['year', 'countries'])
#map the colors of a dict to a dataframe
df['color']= df.countries.map(color_dict)
#we conly need the data from the last year


df

df_last_year = df.dropna()
#map colors and sites
style_map = df_last_year.set_index('sites')['color'].to_dict()
style_map

start = df[df.year==df.year.min()]['sites'].values
end = df[df.year==df.year.max()]['diff'].values
print(start, end)

fig, axes = plt.subplot_mosaic(
    np.fliplr(mosaic), #mirrow the mosaic
    empty_sentinel=0, # no subplots if value = 0
    gridspec_kw={
        "wspace": 0,
        "hspace": 0,
    }
)

for i , (key, ax) in enumerate(reversed(axes.items()), start = 1):
    #print(i)
    if i in style_map:
        fontweight = 'bold'
        color = style_map[i]

    else:
        fontweight = 'light'
        color = 'black'  # or 'gray', etc.
    
    #add numbers
    ax.text(
        0.5,        #pos inside the axes
        0.5, 
        i , 
        ha="center", 
        va="center", 
        fontsize=10, 
        color=color, 
        weight=fontweight
        )

    #styling
    ax.tick_params(length = 0, labelleft = False, labelbottom = False)
    ax.spines[['bottom','right']].set_visible(False)
    ax.spines[['top','left']].set_color(spines_color)

    # Add image if position matches
    if i in flag_map:
        ib = OffsetImage(flag_map[i], zoom=.03)
        ab = AnnotationBbox(
            ib,
            (0,0),
            xybox=(0.5, 1.6),
            alpha=0.5,
            frameon=False,
            )
        ax.add_artist(ab)

    #add the arrows
    x_arr = .5
    y_arr = 2.5
    rad = 0.6
    if i  in set(start):    
        index = list(start).index(i) #Find the index of i in start
        ax.annotate(
            "", 
            xy=(x_arr, y_arr),
            xytext=(end[index] + x_arr , end[index] + y_arr), 
            annotation_clip= False,
            arrowprops=dict(
                arrowstyle='-', 
                connectionstyle=f'arc3,rad={rad}',
                color=df.color.unique()[index], 
                linewidth=2, 
                linestyle='-', 
                antialiased=True
                ))

        #add the labels
        p1 = np.array([x_arr , y_arr])
        p2 = np.array([end[index] + x_arr , end[index] + y_arr])
        mid = (p1 + p2) / 2                             #The center of the chord between the points
        vec = p2 - p1                                   #the direction of the chord
        length = np.linalg.norm(vec)                    #the length of the chord
        perp = np.array([-vec[1], vec[0]]) / length     #the direction fo the arc (upwards, downwards)
        
        arc_height = rad * (length / 2)                 #The offset along the perpendicular, controlled by rad
        arc_mid = mid + arc_height * perp               #The actual midpoint of the arc, i.e. highest point of the curve
        label_x, label_y = arc_mid

        ax.text(
            label_x , 
            label_y + 1, 
            f"+{int(end[index])}", 
            color = df.color.unique()[index],
            weight = "bold",
            fontsize = 10
            ) 


fig