import serial
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d.art3d import Poly3DCollection
from matplotlib.animation import FuncAnimation
import re

# Serial settings (change COM port and baud rate as needed)
SERIAL_PORT = '/dev/ttyACM0'   # e.g., 'COM3' on Windows or '/dev/ttyUSB0' on Linux/Mac
BAUD_RATE = 115200

# Open serial connection
ser = serial.Serial(SERIAL_PORT, BAUD_RATE, timeout=1)

# Parse "roll: xx, pitch: xx, yaw: xx"
def parse_data(line):
    match = re.search(r'roll:\s*(-?\d+\.?\d*),\s*pitch:\s*(-?\d+\.?\d*),\s*yaw:\s*(-?\d+\.?\d*)', line)
    if match:
        roll = float(match.group(1))
        pitch = float(match.group(2))
        yaw = float(match.group(3))
        return roll, pitch, yaw
    return None

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

    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

def create_box():
    l, w, h = 1, 0.5, 0.2
    return 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],
    ])

def get_faces(vertices):
    return [
        [vertices[j] for j in [0,1,2,3]],
        [vertices[j] for j in [4,5,6,7]],
        [vertices[j] for j in [0,1,5,4]],
        [vertices[j] for j in [2,3,7,6]],
        [vertices[j] for j in [1,2,6,5]],
        [vertices[j] for j in [3,0,4,7]],
    ]

# Create plot
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
box = create_box()
c = ['C0'] * 1 + ['C3'] + ['C0'] * 4
collection = Poly3DCollection([], facecolors=c, edgecolors='black', alpha=0.90) #facecolors='skyblue',edgecolors='black', alpha=0.8
ax.add_collection3d(collection)

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.view_init(elev=30, azim=45)

# Animation update function
def update(frame):
    line = ser.readline().decode('utf-8').strip()
    result = parse_data(line)
    if result:
        roll, pitch, yaw = result
        R = euler_to_rotation_matrix(roll, pitch, yaw)
        rotated_box = box @ R.T
        faces = get_faces(rotated_box)
        collection.set_verts(faces)
        ax.set_title(f"Roll={roll:.1f}°, Pitch={pitch:.1f}°, Yaw={yaw:.1f}°")

ani = FuncAnimation(fig, update, interval=50)

plt.tight_layout()
plt.show()