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

# Box and stack parameters (input units can vary)
length = float(input('Input length: '))  # Box length
width = float(input('Input width: '))  # Box width
height = float(input('Input height: '))  # Box height
rows = int(input('Input number of rows per layer: '))  # Number of rows per stack
columns = int(input('Input number of columns per layer: '))  # Number of columns per stack
stacks = int(input('Input number of layer: '))  # Number of stacks placed vertically (reduced for visual clarity)

# Input units (e.g., inches, cm, feet)
unit = str(input('Input unit (e.g inches, cm): '))  # You can change this to any other unit

# Automatically calculate pallet width and length
pallet_width = columns * width
pallet_length = rows * length

# Calculate the max dimension to create a bounding box
max_size = max(pallet_width, pallet_length, stacks * height)

fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')

# Function to draw a single box
def draw_box(ax, x, y, z, width, length, height, color='white', edgecolor='black', annotate=False, unit="inches"):
    # Define the vertices of the box
    vertices = np.array([[x, y, z],
                         [x + width, y, z],
                         [x + width, y + length, z],
                         [x, y + length, z],
                         [x, y, z + height],
                         [x + width, y, z + height],
                         [x + width, y + length, z + height],
                         [x, y + length, z + height]])

    # Define the faces of the box
    faces = [[vertices[j] for j in [0, 1, 2, 3]],  # Bottom face
             [vertices[j] for j in [4, 5, 6, 7]],  # Top face
             [vertices[j] for j in [0, 1, 5, 4]],  # Front face
             [vertices[j] for j in [2, 3, 7, 6]],  # Back face
             [vertices[j] for j in [1, 2, 6, 5]],  # Right face
             [vertices[j] for j in [0, 3, 7, 4]]]  # Left face

    # Draw the box with an edgecolor outline
    ax.add_collection3d(Poly3DCollection(faces, facecolors=color, edgecolors=edgecolor, linewidths=1, alpha=0.9))

    # Annotate the box dimensions only if annotate=True
    if annotate:
        center_x = x + width / 2
        center_y = y + length / 2
        center_z = z + height / 2
        ax.text(center_x, center_y, center_z, f'{width}x{length} {unit}', color='red', ha='center', fontsize=10)

# Draw the boxes in stacks (starting on top of the pallet)
for stack in range(stacks):
    for row in range(rows):
        for col in range(columns):
            x = col * width
            y = row * length
            z = stack * height  # Start drawing above the pallet
            draw_box(ax, x, y, z, width, length, height, unit=unit)

# Manually set equal limits for all axes based on the max size (bounding box)
ax.set_xlim([0, max_size])
ax.set_ylim([0, max_size])
ax.set_zlim([0, max_size])

# Set axis labels with the inputted units
ax.set_xlabel(f'Width ({unit})')
ax.set_ylabel(f'Length ({unit})')
ax.set_zlabel(f'Height ({unit})')
plt.title(f'3D Stacked Boxes on a Pallet (Dimensions in {unit})')

# Show the plot
plt.show()