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

def euler_to_rotation_matrix(roll, pitch, yaw):
    # Convert to radians
    roll = np.radians(roll)
    pitch = np.radians(pitch)
    yaw = np.radians(yaw)

    # Rotation matrices
    Rx = np.array([
        [1, 0, 0],
        [0, np.cos(roll), -np.sin(roll)],
        [0, np.sin(roll), np.cos(roll)]
    ])

    Ry = np.array([
        [np.cos(pitch), 0, np.sin(pitch)],
        [0, 1, 0],
        [-np.sin(pitch), 0, np.cos(pitch)]
    ])

    Rz = np.array([
        [np.cos(yaw), -np.sin(yaw), 0],
        [np.sin(yaw), np.cos(yaw), 0],
        [0, 0, 1]
    ])

    return Rz @ Ry @ Rx  # ZYX order

def create_box():
    # Define 8 corners of a box centered at origin
    l, w, h = 1, 0.5, 0.2  # length, width, height
    corners = np.array([
        [-l/2, -w/20, -h/2],
        [ l/2, -w/2, -h/2],
        [ l/2,  w/2, -h/2],
        [-l/2,  w/20, -h/2],
        [-l/2, -w/20,  h/2],
        [ l/2, -w/2,  h/2],
        [ l/2,  w/2,  h/2],
        [-l/2,  w/20,  h/2],
    ])
    return corners

def get_faces(vertices):
    # Faces defined by vertices index
    faces = [
        [vertices[j] for j in [0,1,2,3]],  # Bottom
        [vertices[j] for j in [4,5,6,7]],  # Top
        [vertices[j] for j in [0,1,5,4]],  # Front
        [vertices[j] for j in [2,3,7,6]],  # Back
        [vertices[j] for j in [1,2,6,5]],  # Right
        [vertices[j] for j in [3,0,4,7]],  # Left
    ]
    return faces

def plot_box(roll, pitch, yaw):
    box = create_box()
    R = euler_to_rotation_matrix(roll, pitch, yaw)
    
    # Rotate box
    rotated_box = box @ R.T  # Apply rotation
    
    faces = get_faces(rotated_box)
    
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')

    # Draw box
    c = ['C0'] * 1 + ['C3'] + ['C0'] * 4
    ax.add_collection3d(Poly3DCollection(faces, facecolors=c, edgecolors='black', alpha=0.90))

    # Axis setup
    ax.set_xlim([-1, 1])
    ax.set_ylim([-1, 1])
    ax.set_zlim([-1, 1])
    ax.set_xlabel('X')
    ax.set_ylabel('Y')
    ax.set_zlabel('Z')
    ax.set_title(f"3D Box Orientation\nRoll={roll}°, Pitch={pitch}°, Yaw={yaw}°")
    ax.view_init(elev=30, azim=45)
    plt.tight_layout()
    plt.show()

# Example usage
if __name__ == "__main__":
    roll = float(input("Enter roll (degrees): "))
    pitch = float(input("Enter pitch (degrees): "))
    yaw = float(input("Enter yaw (degrees): "))
    plot_box(roll, pitch, yaw)