from matplotlib import pyplot as plt
from matplotlib.offsetbox import OffsetImage, AnnotationBbox
from mpl_toolkits.axes_grid1.inset_locator import inset_axes  #to align axes
import pandas as pd
import numpy as np

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

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 = df.pivot_table( index= 'countries', columns = 'year', values = 'sites')
df = df.astype(int)
df.loc['Avg.'] = df.mean(numeric_only=True).astype(int)
df['Change'] = (df.iloc[:, -1] - df.iloc[:, 0]).astype(int)
df['in %'] = (df['Change']/df.iloc[:, 0]*100).round(1)
df.columns.name = None #dont name the index
df = df.reset_index()
df = df.sort_index( ascending=False )
df

#Define the table
rows = len(df.countries)
cols = len(df.columns[1:])

#Position of the values in the table
value_positions = list(np.arange(0, cols, 1 )) 
print(value_positions)

# Max of each column 
numeric_cols = df.select_dtypes(include='number').columns
max_row = df[numeric_cols].max()

fig, ax = plt.subplots(figsize=(8,6))

ax.set_ylim(-1, rows + 1)
ax.set_xlim(-1, cols + .5)

cell_height_offset= .1   #move the cells up to have space for the bars

# Add table's values
for i in range(rows):
    for j, column in enumerate(df.columns[1:]):
        value = df[column].iloc[i]
        country = df['countries'].iloc[i]  # Get country for this row
 
        # Create inset axes positioned based on the data coordinates
        bar_ax = inset_axes(
            ax,
            width=1,
            height=.1,
            loc='lower left',  #anchor it on the lower left
            bbox_to_anchor=(value_positions[j] , i - cell_height_offset *2),
            bbox_transform=ax.transData,
            borderpad=0
            )

        bar_color = color_dict.get(country, "#CCCCCC")

        #add the bar
        bar_ax.barh(
            y = 0, 
            width = value,
            color = bar_color
            )

        bar_ax.set(xlim = (0, max_row[column]))
        bar_ax.set_axis_off()

        # if it's the 4th column add '%'
        if j == 3:
            # Only show decimal if it's not .0
            if value == int(value):
                display_value = f"{int(value)}%"    
            else:
                display_value = f"{value:.1f}%"     
        else:
            display_value = int(value) if value == int(value) else value


        ax.annotate(
            xy=(value_positions[j] , i + cell_height_offset), 
            weight = "bold", 
            size= 10,
            text=display_value,
            va='center',
            ha = 'left'
        )
        #add column headers
        ax.annotate(
            xy = (value_positions[j],rows ), 
            text = column, 
            size = 14,
            color = "#818F9A", 
            weight = "light"
        )
    
    # Add dividing lines
    ax.plot(
        [-1, cols ],
        [i-.5, i- .5],
        ls='-',
        lw='.5',
        c='lightgrey', 
        clip_on = False
        )
    ax.plot(
        [-1, cols ],
        [3.5, 3.5], 
        lw='.5',
        c='lightgrey'
        )

x_row_header = -0.5 
for im, values in zip(img, range(1,cols +1)):
    image_box =  OffsetImage(im, zoom = 0.05) #container for the image
    ab = AnnotationBbox(
        image_box, 
        (x_row_header, values),
        xycoords='data',
        frameon = False,
        clip_on = False
        )
    ax.add_artist(ab)

ax.annotate(
    "Avg.", 
    xy= (x_row_header , 0), 
    size= 16, 
    color = "#B4BDC3" ,
    ha = "center",
    va = "center"
    )

ax.set_axis_off()
fig