#!/usr/bin/env python

import rospy
from geometry_msgs.msg import Quaternion
from sensor_msgs.msg import Imu, MagneticField
from std_msgs.msg import Header
import serial
import math
import json

class IMUPublisher:
    def __init__(self):
        rospy.init_node('arduino_serial_reader')
        self.serial_port = '/dev/ttyACM0'
        self.baud_rate = 115200
        
        self.publisher_mag = rospy.Publisher('/imu/mag', MagneticField, queue_size=10)
        self.publisher_imu = rospy.Publisher('/imu/imu', Imu, queue_size=10)
        
        self.timer = rospy.Timer(rospy.Duration(1/100), self.publish_magnet)
        self.timer = rospy.Timer(rospy.Duration(1/100), self.publish_imu)
        self.timer = rospy.Timer(rospy.Duration(1/100), self.read_serial_data)
        
        self.mag = MagneticField()
        self.imu = Imu() 
        
        self.qx = 0.0
        self.qy = 0.0
        self.qz = 0.0
        self.qw = 0.0      
        self.euler_roll  = 0.0
        self.euler_pitch = 0.0
        self.euler_yaw   = 0.0
        self.angular_vel_x  = 0.0
        self.angular_vel_y  = 0.0
        self.angular_vel_z  = 0.0
        self.linear_accel_x = 0.0
        self.linear_accel_y = 0.0
        self.linear_accel_z = 0.0
        self.magnetometer_x = 0.0
        self.magnetometer_y = 0.0
        self.magnetometer_z = 0.0

        try:
            self.ser = serial.Serial(self.serial_port, self.baud_rate)
            rospy.loginfo("Serial port opened: %s", self.serial_port)
        except serial.SerialException:
            rospy.logerr("Failed to open serial port. Make sure it's correct and not already in use.")
            return

    def read_serial_data(self, event):
        if self.ser.is_open:
            try:
                received_data = self.ser.read_until(b'\n').decode('utf-8')
                try:
                    data_dict = json.loads(received_data)
                    self.process_data(data_dict)
                except json.JSONDecodeError:
                    rospy.logerr("Error decoding JSON data.")
            except UnicodeDecodeError:
                rospy.logerr("Error decoding received data.")
        else:
            rospy.logwarn("Serial port is not open.")

    def process_data(self, data_dict):
        if "A" in data_dict:
            euler_data = data_dict["A"]
            self.euler_roll  = euler_data[0]
            self.euler_pitch = euler_data[1]
            self.euler_yaw   = euler_data[2]
            rospy.loginfo("Euler Angle: euler_roll=%.2f, euler_pitch=%.2f, euler_yaw=%.2f", self.euler_roll, self.euler_pitch, self.euler_yaw)
                            
        if "G" in data_dict:
            angular_vel_data = data_dict["G"]
            self.angular_vel_x = angular_vel_data[0]
            self.angular_vel_y = angular_vel_data[1]
            self.angular_vel_z = angular_vel_data[2]
            rospy.loginfo("Angular Velocity: X=%.2f, Y=%.2f, Z=%.2f", self.angular_vel_x, self.angular_vel_y, self.angular_vel_z)

        if "L" in data_dict:
            linear_accel_data = data_dict["L"]
            self.linear_accel_x = linear_accel_data[0]
            self.linear_accel_y = linear_accel_data[1]
            self.linear_accel_z = linear_accel_data[2]
            rospy.loginfo("Linear Acceleration: X=%.2f, Y=%.2f, Z=%.2f", self.linear_accel_x, self.linear_accel_y, self.linear_accel_z)

        if "M" in data_dict:
            magnetometer_data = data_dict["M"]
            self.magnetometer_x = magnetometer_data[0]
            self.magnetometer_y = magnetometer_data[1]
            self.magnetometer_z = magnetometer_data[2]
            rospy.loginfo("Magnetometer: X=%.2f, Y=%.2f, Z=%.2f", self.magnetometer_x, self.magnetometer_y, self.magnetometer_z)
            
    def euler_to_quaternion(self):
        sr = math.sin(math.radians(self.euler_roll) /2)
        cr = math.cos(math.radians(self.euler_roll) /2)
        sp = math.sin(math.radians(self.euler_pitch)/2)
        cp = math.cos(math.radians(self.euler_pitch)/2)
        sy = math.sin(math.radians(self.euler_yaw)  /2)
        cy = math.cos(math.radians(self.euler_yaw)  /2)

        self.qw = cy * cp * cr + sy * sp * sr
        self.qx = cy * cp * sr - sy * sp * cr
        self.qy = sy * cp * sr + cy * sp * cr
        self.qz = sy * cp * cr - cy * sp * sr
                                                
    def publish_imu(self, event):
        self.imu.header = Header()
        self.imu.header.frame_id = "base_link"
        self.imu.header.stamp = rospy.Time.now()
        
        self.euler_to_quaternion()
        self.imu.orientation.x = self.qx
        self.imu.orientation.y = self.qy
        self.imu.orientation.z = self.qz
        self.imu.orientation.w = self.qw
        self.imu.orientation_covariance = [
            0.0159, 0.0, 0.0,
            0.0, 0.0159, 0.0,
            0.0, 0.0, 0.0159]
        rospy.loginfo("Orientation: qx=%.2f, qy=%.2f, qz=%.2f, qw=%.2f", self.imu.orientation.x, self.imu.orientation.y, self.imu.orientation.z, self.imu.orientation.w)

        self.imu.angular_velocity.x = float(self.angular_vel_x)
        self.imu.angular_velocity.y = float(self.angular_vel_y)
        self.imu.angular_velocity.z = float(self.angular_vel_z)
        self.imu.angular_velocity_covariance = [
            0.04, 0.0, 0.0,
            0.0, 0.04, 0.0,
            0.0, 0.0, 0.04]
        rospy.loginfo("Angular Velocity: X=%.2f, Y=%.2f, Z=%.2f", self.imu.angular_velocity.x, self.imu.angular_velocity.y, self.imu.angular_velocity.z)

        self.imu.linear_acceleration.x = float(self.linear_accel_x)
        self.imu.linear_acceleration.y = float(self.linear_accel_y)
        self.imu.linear_acceleration.z = float(self.linear_accel_z)
        self.imu.linear_acceleration_covariance = [
            0.017, 0.0, 0.0,
            0.0, 0.017, 0.0,
            0.0, 0.0, 0.017]
        rospy.loginfo("Linear Acceleration: X=%.2f, Y=%.2f, Z=%.2f", self.imu.linear_acceleration.x, self.imu.linear_acceleration.y, self.imu.linear_acceleration.z)
                                                   
        self.publisher_imu.publish(self.imu)
        
    def publish_magnet(self, event):
        self.mag.header = Header()
        self.mag.header.frame_id = "base_link"
        self.mag.header.stamp = rospy.Time.now()
        
        self.mag.magnetic_field.x = float(self.magnetometer_x)
        self.mag.magnetic_field.y = float(self.magnetometer_y)
        self.mag.magnetic_field.z = float(self.magnetometer_z)
    
        self.publisher_mag.publish(self.mag)

def main():
    node = IMUPublisher()
    rospy.spin()

if __name__ == '__main__':
    main()

