import matplotlib.pyplot as plt
from matplotlib.patches import FancyBboxPatch
from matplotlib.offsetbox import OffsetImage, AnnotationBbox
from matplotlib.colors import LinearSegmentedColormap
import pandas as pd
import numpy as np

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

xy_ticklabel_color, label_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 = df.sort_values([ 'countries'], ascending=False ).reset_index(drop=True)
df['year_lbl'] ="'"+df['year'].astype(str).str[-2:].astype(str)
df['pct_change'] = (df.groupby('countries', sort=False)['sites'].apply(lambda x: x.pct_change()).to_numpy()).round(3)
#map the colors of a dict to a dataframe
df['color']= df.year.map(color_dict)
df

fig, axes = plt.subplots(ncols = df.countries.nunique(), nrows = 1, figsize=(8,5),sharex=True, sharey=True, facecolor = "#FFFFFF", zorder= 1)
fig.tight_layout(pad=1.0)


for ax, (i,(country, group)) in zip(axes.ravel(), enumerate(df.groupby("countries"))):
    colors = group.color
    x = range(len(df.year.unique()))

    ax.bar(
        x, 
        height = group.sites, 
        width = bar_width,
        color = group.color)

    #add the multicolor lines
    multiColorLine(0 - bar_width / 2, 1 + bar_width / 2, -0.5, -0.5, 900, 0.01, colors.tolist(), ax)    
    multiColorLine(0 - bar_width / 2, 1 + bar_width / 2, group.sites.iloc[0] + line_offset , group.sites.iloc[1] + line_offset, 900, 0.01, colors.tolist(), ax)

    #add the grand totals at the top of the bars
    for bar, site in zip(ax.patches, group.sites):
            ax.text(
            bar.get_x() + bar.get_width() / 2, 
            round(bar.get_height())+0.5,  #height
            site, 
            ha="center", 
            va="bottom", 
            size= 12,
            color = label_color, 
            weight= "bold"
        )

    #add the image
    image_box =  OffsetImage(img[i], zoom = 0.05) #container for the image
    ab = AnnotationBbox(
        image_box, 
        (0.5,1),  
        xybox= (0.5, -4), 
        frameon = False) # the x coordinate is the year axis
    ax.add_artist(ab)

    value = (group['pct_change'].iloc[1])*100

    if value == int(value):
        display_value = f"{int(value)}%"    
    else:
        display_value = f"{value:.1f}%" 

    #add data annotations
    ax.annotate(
        f'\u25B2\n{display_value}', 
        xy= (.5, .2),
        size= 12, 
        color= label_color, 
        weight = "bold", 
        va = "center", 
        ha = "center", 
        xycoords="axes fraction"
        )


    #add styling
    #round the edges of the bars
    new_patches = []
    for patch in reversed(ax.patches):
        bb = patch.get_bbox()
        color = patch.get_facecolor()
        p_bbox = FancyBboxPatch(
            (bb.xmin, bb.ymin),
            abs(bb.width), 
            abs(bb.height),
            boxstyle="round,pad=.005,rounding_size=.2",
            ec="none", 
            fc=color,
            mutation_aspect=0.2
            )
        patch.remove()
        new_patches.append(p_bbox)

    for patch in new_patches:
        ax.add_patch(patch)
  
    ax.set_frame_on(False)
    ax.set_xlabel(
        country, 
        size=14, 
        color = label_color,
        labelpad = 58
        )


    ax.xaxis.set_ticks(
        x, 
        labels =group.year_lbl, 
        color= label_color
        )

    ax.tick_params(
        axis='both', 
        which='major',
        length=0, 
        labelsize=11,
        colors= label_color, 
        labelleft = False
        )