I am trying to create a figure to explain something in a report I am writing.

My figure is just a bynch of cubes stacked axially.

I want to add a curly bracket ‘{‘ next to ech level in my stacked cube figure and add text next to the bracket to label that particular section.

I made a code that makes the cubes but I dont know how to add the brackets.

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

def draw_cube(ax, origin, size, height, label, dotted = False):
    # Vertices of a cube
    x, y, z = origin
    vertices = [
        [x, y, z],
        [x + size, y, z],
        [x + size, y + size, z],
        [x, y + size, z],
        [x, y, z + height],
        [x + size, y, z + height],
        [x + size, y + size, z + height],
        [x, y + size, z + height]
    ]
    
    # Define the 6 faces of the cube
    faces = [
        [vertices[0], vertices[1], vertices[2], vertices[3]],
        [vertices[4], vertices[5], vertices[6], vertices[7]], 
        [vertices[0], vertices[1], vertices[5], vertices[4]], 
        [vertices[2], vertices[3], vertices[7], vertices[6]], 
        [vertices[1], vertices[2], vertices[6], vertices[5]],
        [vertices[4], vertices[7], vertices[3], vertices[0]]
    ]
    linestyle = 'dotted' if dotted else '-'
    # Draw the faces
    ax.add_collection3d(Poly3DCollection(faces, linewidths=1, edgecolors='k', facecolor = 'b',alpha=.25, linestyle = linestyle))


def main():
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d', elev= 10, azim = 15)

    # cubes for section 1
    draw_cube(ax, origin=(0, 0, 0), size=1, height = 1, dotted = False, label= "$U_{j}$")

    # cubes for section 2
    draw_cube(ax, origin=(0, 0, 1), size=0.5, height = 3, dotted = False, label = "$U_{j+1}$")
    draw_cube(ax, origin=(0, 0.5, 1), size=0.5, height = 3, dotted = False, label = "$U_{j+1}$")
    draw_cube(ax, origin=(0.5, 0, 1), size=0.5, height = 3, dotted = False, label = "$U_{j+1}$")
    draw_cube(ax, origin=(0.5, 0.5, 1), size=0.5, height = 3, dotted = False, label = "$U_{j+1}$")

    #cubes for section 3
    draw_cube(ax, origin=(0, 0, 4), size=1, height = 2, dotted = False, label= "$U_{j}$")

    #draw_cube(ax, origin=(0, 0, 25), size=1, height = 50, dotted = False)

    # Set the aspect ratio to be equal
    ax.set_aspect('auto')

    # Set labels
    ax.set_xlabel('X')
    ax.set_ylabel('Y')
    ax.set_zlabel('Z')

    ax.axis('off')
    # Set limits
    ax.set_xlim([0, 1])
    ax.set_ylim([0, 1.5])
    ax.set_zlim([0, 4.1])

    # Display the plot
    plt.show()
if __name__ == "__main__":
    main()

This code produces the following figure

I dont know how to add curly brackets to a 3d plot