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

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

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['ctry_code'] = df.countries.map(code_dict)
df['color']=df.countries.map(color_dict)
df["pct_change"] = (
    df.groupby("countries", sort=False)["sites"]
    .apply(lambda x: x.pct_change())
    .to_numpy()
    .round(3)
    * 100
)
df = df.sort_values(['sites' ], ascending=False ).reset_index(drop=True)
df

fig = plt.figure(figsize=(6,6))
ax = fig.add_subplot(1, 1, 1, projection="3d",computed_zorder=False)

ax.set(xlim = [1,8], ylim = [1,6], zlim = [0,16])

y = [2,4]*3
x = [2,2,4,4,6,6]
z = [0]*6

width = [1]*6 #width
depth = [1]*6 #depth

height=[] #height
for country in df.countries.unique():
    sites = df[df.countries == country].sort_values("year", ascending =True)["sites"].values
    height.extend(sites.tolist())


ax.bar3d(x, y, z, width, depth, height, color=(0, 0, 1, 0),  ec= "lightgrey",lw=0.5, zorder=0) #ec= "#F0F2F3"

blocks = get_bar_vertices(x, y, z, width, depth, height)
print(blocks)

pct_changes = df['pct_change'].dropna().unique()
pct_changes 

colors = df['color'].unique()
pct_changes = df['pct_change'].dropna().unique()

for i in range(0, len(blocks), 2):
    current_color = colors[(i // 2) % len(colors)]

    x_left, x_right, y_front, y_back, z_bottom, z_top = blocks[i]  #front block
    X_Left, X_Right, Y_Front, Y_Back, Z_Bottom, Z_Top = blocks[i + 1] #black block

    block = np.array([

        #first surface
        [[x_left, y_front, z_top],
         [x_right, y_front, z_top],
         [x_right, y_back, z_top],
         [x_left, y_back, z_top]],

        #middle surface
        [[x_left, y_back, z_top],
         [x_right, y_back, z_top],
         [X_Right, Y_Front, Z_Top],
         [X_Left, Y_Front, Z_Top]],

        #last surface
        [[X_Left, Y_Front, Z_Top],
         [X_Right, Y_Front, Z_Top],
         [X_Right, Y_Back, Z_Top],
         [X_Left, Y_Back, Z_Top]]
    ])

    pc = Poly3DCollection(block, facecolors=current_color, shade=True, alpha=1 )
    ax.add_collection3d(pc)

    #Add text middle surface
    xc = (x_left + x_right + X_Left + X_Right) / 4
    yc = (y_back + Y_Front) / 2
    zc = (z_top + Z_Top) / 2

    pct = pct_changes[i // 2]

    ax.text(
        xc, yc, zc,
        f"+{int(pct)}%",
        ha="center", va="center", color="white", size=8, zorder=3, weight = "bold"
    )


ax.xaxis._axinfo["grid"].update({"linewidth":0, "color" : "w"}) #color gridline
ax.xaxis.set_ticks(np.arange(2, 8, 2), labels = ["SE","DK","NO"])
ax.set_yticks([])
ax.set_zticks([])

label_colors  = [ "#5375D4",  "#A54836", "#2B314D"]
for xtick, color in zip(ax.get_xticklabels(), label_colors):
    xtick.set_color(color)

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")    

# Transparent spines
ax.xaxis.line.set_color("w")   
ax.yaxis.line.set_color("w")
ax.zaxis.line.set_color("w")
fig

for i, block in enumerate(blocks):
    x1, x2, y1, y2, z1, z2 = block
    #calculate the center of the surfaces
    xc = (x1 + x2) / 2
    yc = (y1 + y2) / 2

    print(i, xc,yc,z2, xc + 0.2, yc + 0.2, z2)
    ax.text(
        xc, yc, z2,
        f"{z2}",
        ha="center", va="center", color="white", zorder=3
    )

fig