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

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

xy_ticklabel_color, flag_text_color, grid_color, datalabels_color = (
    "#101628",
    "#101628",
    "#E3E6E8",
    "#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["pct_tot"] = df["sites"] / df.groupby("year")["sites"].transform("sum")
df["sub_total"] = df.groupby("countries")["sites"].transform("sum")
df["ctry_code"] = df.countries.map(code_dict)

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

# map the colors of a dict to a dataframe
df["color"] = df.countries.map(color_dict)
df

years = df.year.unique()
codes = df.ctry_code.unique()

#load images
img = [
    plt.imread("../flags/de-rd.png"),
    plt.imread("../flags/no-rd.png"),
    plt.imread("../flags/sw-rd.png"),
] * 3

# get first and last color of the series to color the end bars
color_end_bars = df.color[:1].to_list() + df.color[-1:].to_list()

groups  = df.groupby("color", sort = False)['pct_tot'].apply(np.array)
groups

fig, ax = plt.subplots(figsize=(10, 4), facecolor="#FFFFFF")

# add the stacked bars
bottom = np.zeros(len(years))
for color, y in groups.items():
    ax.barh(
        range(len(years)), 
        y, 
        left = bottom, 
        height = 0.42, 
        zorder = 1, 
        color = color)
    bottom += y

# add the end of the bars
ax.scatter(
    np.array([0, 1] * 2),
    np.sort([0, 1] * 2),
    marker = "o",
    s = 1650,
    color = color_end_bars * 2,
    zorder = 2,
)

ax.set(xlim=([-0.11, 1.1]), ylim=(-0.4, 1.8))

fig

# add data labels
for bar, pct in zip(ax.patches, df.pct_tot):
    
    xbar = bar.get_x() + bar.get_width() / 2
    ybar = bar.get_height() / 2 + bar.get_y()

    ax.text(
        xbar,
        ybar,
        f"{pct:.0%}",
        ha="center",
        va="center",
        color=datalabels_color,
        size=16,
    )
fig

# add flags only on one axis
for bar, im, code in zip(ax.patches[1::2], img, codes):
    xbar = bar.get_x() + bar.get_width() / 2
    ybar = bar.get_height() / 2 + bar.get_y()

    image_box = OffsetImage(im, zoom=0.06)  # container for the image
    ab = AnnotationBbox(
        image_box, 
        (xbar, ybar + 0.8), 
        frameon=False)

    ax.text(
        xbar, 
        ybar + 0.45, 
        code, 
        size = 12, 
        ha = "center", 
        color = flag_text_color
        )

    ax.add_artist(ab)

fig

ax.plot(
    [0, 1],
    np.array([[0, 0], [1, 1]]).T,
    linestyle="-",
    lw=50,
    color="w",
    path_effects=[pe.Stroke(
        linewidth=52, 
        foreground=grid_color), 
        pe.Normal()],
    solid_capstyle="round",
    zorder=0,
)

fig

# add straight line
ax.hlines(
    y=[0, 1],             # list of y positions
    xmin=-1, 
    xmax=0.5,
    color=grid_color,
    zorder=-1
)

# add years
ax.yaxis.set_ticks(
    range(len(years)), 
    labels = years)

fig

ax.tick_params(
    axis='y', 
    which='major',
    length=0, 
    labelsize=14,
    colors= xy_ticklabel_color,
    )
ax.set_frame_on(False)

ax.set_xticks([])

fig