import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d.art3d import Poly3DCollection
import numpy as np
import pandas as pd

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

xy_ticklabel_color,  line_colors, datalabels_color ='#757C85', "#757C85", "#757C85"

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['ctry_code'] = df.countries.map(code_dict)
df['year_lbl'] ="'"+df['year'].astype(str).str[-2:].astype(str)
# custom sort
sort_order_dict = {"Denmark": 3, "Sweden": 1, "Norway": 2, 2004: 4, 2022: 5}
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

import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d.art3d import Poly3DCollection

fig, ax = plt.subplots(figsize=(5,8), subplot_kw={'projection':'3d'})

# Common dimensions
depth = 0.6
bar_width = 0.2
y_start = 0

ax.set(xlim = [0,1.1],ylim=[0, 1], zlim=[0, df.sites.max()])

# Variable dimensions
x_positions = np.linspace(0, .8, len(df.countries.unique()))

grouped_sites = df.groupby("color", sort=False)['sites'].apply(list)

for (color, site_pair), x_start in zip(grouped_sites.items(), x_positions):
    height_front, height_back = site_pair

    create_inclined_bar(
        x_start=x_start,
        y_start=y_start,
        width=bar_width,
        depth=depth,
        height_front=height_front,
        height_back=height_back,
        color=color
    )

ax.view_init(
    elev=30,    # looking from 30 degrees above the x-y plane (higher/ lower above the plot)
    azim=-50,   # rotates camera around z-axis (left/ right around the plot)
    roll=0)    # tilts the view

# midpoint of each block in the x coordinates
x_midpoint = x_positions + bar_width / 2
xs = np.repeat(x_midpoint, 2)

# y coordinates with offset
offset = .1
ys = [y_start - offset, y_start + depth + offset] * (len(xs) // 2)

z_positions = df['sites'].tolist() 

# alternate sites values for 2022 with zeros
zs =  [0 if i % 2 == 0 else z for i, z in enumerate(z_positions)] 

# texts: same as zs
texts = z_positions

# Print results
print("xs =", xs)
print("ys =", ys)
print("zs =", zs)
print("texts =", texts)

#add data labels
for  text, x, y, z in zip( texts, xs, ys, zs):
    ax.text(
        x, 
        y, 
        z, 
        text,
        va="center", 
        color = datalabels_color,
        ha="right", 
        clip_on=True
    )

fig

# Define axis line start and end points
axis_lines = {
    'start': {'x': [0, 1.1], 'y': [depth, depth], 'z': [0, 0]},
    'end': {'x': [0, 1.1], 'y': [0, 0], 'z': [0, 0]},
}

# Plot axis lines
for axis, coords in axis_lines.items():
    ax.plot(
        coords['x'],     # x-coordinates
        coords['y'],         # y-coordinates
        zs = coords['z'],      # z-coordinates
        color = line_colors,
        lw = 1
    )


fig

#add country labels
ax.xaxis.set_ticks(x_midpoint -.1, labels = df.ctry_code.unique())

ax.tick_params(
    axis='x', 
    which='major', 
    length=0, 
    labelsize=11,
    pad =10
    )


for xtick, color in zip(ax.get_xticklabels(), df.color.unique()):
    xtick.set_color(color)

#add year labels
ax.yaxis.set_ticks(
    [0 + offset /2, depth + offset /2], 
    labels = df.year_lbl.unique(), 
    color = xy_ticklabel_color
    )

ax.zaxis.set_ticks([])

fig

for axis in [ax.xaxis, ax.yaxis, ax.zaxis]:
    axis._axinfo['tick']['inward_factor'] = 0
    axis._axinfo['tick']['outward_factor'] = 0      #remove ticks
    axis.set_pane_color("w")           # Make panes transparent
    axis.line.set_linewidth(0)

ax.grid(False)

fig