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


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

xy_ticklabel_color, label_color, grid_color, datalabels_color ='#757C85',"#444E56", "#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['pct_change'] = (df.groupby('countries', sort=False)['sites'].apply(lambda x: x.pct_change()).to_numpy()).round(3)
df["color"] = df.set_index(["year", "countries"]).index.map(color_dict_bar.get)
df['year_lbl'] = "'"+df['year'].astype(str).str[-2:].astype(str)
df = df.sort_values(["year"], ascending=False)
df

fig, ax = plt.subplots(figsize=(8,8), facecolor = "#FFFFFF" )
from matplotlib.colors import LinearSegmentedColormap
ax.set_prop_cycle(color =df.color.unique())  #to control which colors matplotlib cycles through

for i, (x, y, (country, group)) in enumerate(zip(x_coords, y_coords, df.groupby("countries"))):
    x_new = np.linspace(min(x), max(x), 500)
    
    f = interp1d(x, y, kind='quadratic')
    y_smooth=f(x_new)

    #create the gradient lines
    cmap = LinearSegmentedColormap.from_list("", group.color.tolist())
    norm = plt.Normalize(x_new.min(), x_new.max())
    line_colors = cmap(norm(x_new))
    
    ax.scatter(
        x_new,
        y_smooth,
        lw=.2,
        color = line_colors)


    value = (group['pct_change'].iloc[0])*100
        
    if value == int(value):
        display_value = f"+{int(value)}%"    
    else:
        display_value = f"+{value:.1f}%" 
        
    #add data labels
    ax.annotate( 
        display_value, 
        xy=(max(x_new) + .3, max(y_smooth)),
        size = 16,
        color = label_color, 
        weight= "bold", 
        annotation_clip=False
        )

positionx = [df['pct_change'].max() +.2]*3 + [-.1,-.15, -.20]
positiony = df['pct_change'].dropna().tolist() + [0]*3  

line_y = [0.05, 0.95]

for i, (posx, posy) in enumerate(zip(positionx, positiony)):  

    image_box =  OffsetImage(img[i], zoom = 0.05) #container for the image
    ab = AnnotationBbox(
        image_box, 
        (1,1), 
        xybox=(posx, posy), 
        frameon = False)
    ax.add_artist(ab)

fig

line_y = [0.05, 0.95]
line_x = [0, df['pct_change'].max()]
year_labels = df.year_lbl.unique().tolist()

for i in range(2):
    x_line = line_x[i]
    y_line = line_y[i]
    print(x_line,y_line)
    #add the year labels lines
    ax.axvline(
        x=x_line, 
        ymin=0, 
        ymax=y_line,
        lw=1,
        zorder=0, 
        color= "#DFE3E5"
        )

        #add the year labels
    ax.annotate( 
        year_labels[i], 
        xy=(line_x[i], -.15),
        size = 16,
        color = "#9EA6AC", 
        ha = "center",
        va = "center",
        annotation_clip=False
        )



ax.set_axis_off()
fig