#!/usr/bin/env python

import smbus
import rospy
import string
import math
import time
import sys

from tf.transformations import quaternion_from_euler
from dynamic_reconfigure.server import Server
from rasp_imu_hat_6dof.cfg import imuConfig
from sensor_msgs.msg import Imu


degrees2rad = math.pi/180.0
imu_yaw_calibration = 0.0

# Callback for dynamic reconfigure requests
def reconfig_callback(config, level):
    global imu_yaw_calibration
    imu_yaw_calibration = config['yaw_calibration']
    rospy.loginfo("Set imu_yaw_calibration to %d" % (imu_yaw_calibration))
    return config


rospy.init_node("imu_node")

#Get DIY param
topic_name = rospy.get_param('~pub_topic', 'imu')
link_name  = rospy.get_param('~link_name', 'base_imu_link')
pub_hz     = rospy.get_param('~pub_hz', '10')


pub = rospy.Publisher(topic_name, Imu, queue_size=1)
srv = Server(imuConfig, reconfig_callback)  # define dynamic_reconfigure callback
imuMsg = Imu()

# Orientation covariance estimation:
# Observed orientation noise: 0.3 degrees in x, y, 0.6 degrees in z
# Accelerometer non-linearity: 0.2% of 4G => 0.008G. This could cause
# static roll/pitch error of 0.8%, owing to gravity orientation sensing
# error => 2.8 degrees, or 0.05 radians. i.e. variance in roll/pitch: 0.0025
# so set all covariances the same.
imuMsg.orientation_covariance = [
0.0025 , 0 , 0,
0, 0.0025, 0,
0, 0, 0.0025
]

# Angular velocity covariance estimation:
# Observed gyro noise: 4 counts => 0.28 degrees/sec
# nonlinearity spec: 0.2% of full scale => 8 degrees/sec = 0.14 rad/sec
# Choosing the larger (0.14) as std dev, variance = 0.14^2 ~= 0.02
imuMsg.angular_velocity_covariance = [
0.02, 0 , 0,
0 , 0.02, 0,
0 , 0 , 0.02
]

# linear acceleration covariance estimation:
# observed acceleration noise: 5 counts => 20milli-G's ~= 0.2m/s^2
# nonliniarity spec: 0.5% of full scale => 0.2m/s^2
# Choosing 0.2 as std dev, variance = 0.2^2 = 0.04
imuMsg.linear_acceleration_covariance = [
0.04 , 0 , 0,
0 , 0.04, 0,
0 , 0 , 0.04
]

seq=0
accel_factor = 9.806/256.0    # sensor reports accel as 256.0 = 1G (9.8m/s^2). Convert to m/s^2.

class MyIMU(object):
    def __init__(self, addr):
        self.addr = addr
        self.i2c = smbus.SMBus(1)

    def get_YPRAG(self):
        try:
            roll_tmp  = self.i2c.read_i2c_block_data(self.addr, 0x3d, 2)
            pitch_tmp = self.i2c.read_i2c_block_data(self.addr, 0x3e, 2)
            yaw_tmp   = self.i2c.read_i2c_block_data(self.addr, 0x3f, 2)

            ax_tmp  = self.i2c.read_i2c_block_data(self.addr, 0x34, 2)
            ay_tmp  = self.i2c.read_i2c_block_data(self.addr, 0x35, 2)
            az_tmp  = self.i2c.read_i2c_block_data(self.addr, 0x36, 2)

            gx_tmp  = self.i2c.read_i2c_block_data(self.addr, 0x37, 2)
            gy_tmp  = self.i2c.read_i2c_block_data(self.addr, 0x38, 2)
            gz_tmp  = self.i2c.read_i2c_block_data(self.addr, 0x39, 2)
        except IOError:
            rospy.logerr("Read IMU YPRAG date error !")
        else:
            self.raw_roll  = float(((roll_tmp[1]<<8) |roll_tmp[0])/32768.0*180.0)
            self.raw_pitch = float(((pitch_tmp[1]<<8)|pitch_tmp[0])/32768.0*180.0)
            self.raw_yaw  = (yaw_tmp[1] << 8 | yaw_tmp[0])/32768.0*180

            self.raw_ax  = float(((ax_tmp[1]<<8)|ax_tmp[0])/32768.0*16.0)
            self.raw_ay  = float(((ay_tmp[1]<<8)|ay_tmp[0])/32768.0*16.0)
            self.raw_az  = float(((az_tmp[1]<<8)|az_tmp[0])/32768.0*16.0)

            self.raw_gx  = float(((gx_tmp[1]<<8)|gx_tmp[0])/32768.0*2000.0)
            self.raw_gy  = float(((gy_tmp[1]<<8)|gy_tmp[0])/32768.0*2000.0)
            self.raw_gz  = float(((gz_tmp[1]<<8)|gz_tmp[0])/32768.0*2000.0)

    def get_quatern(self):
        try:
            q0 = self.i2c.read_i2c_block_data(self.addr, 0x51, 2)
            q1 = self.i2c.read_i2c_block_data(self.addr, 0x52, 2)
            q2 = self.i2c.read_i2c_block_data(self.addr, 0x53, 2)
            q3 = self.i2c.read_i2c_block_data(self.addr, 0x54, 2)
        except IOError:
            rospy.logerr("Read IMU quaternion date error !")
        else:
            self.raw_q0 = float(((q0[1]<<8)|q0[0])/32768.0)
            self.raw_q1 = float(((q1[1]<<8)|q1[0])/32768.0)
            self.raw_q2 = float(((q2[1]<<8)|q2[0])/32768.0)
            self.raw_q3 = float(((q3[1]<<8)|q3[0])/32768.0)


myIMU = MyIMU(0x50)
rate = rospy.Rate(pub_hz)

while not rospy.is_shutdown():
    myIMU.get_YPRAG()

    yaw_deg = float(myIMU.raw_yaw)
    yaw_deg = yaw_deg + imu_yaw_calibration
    if yaw_deg >= 180.0:
        yaw_deg -= 360.0

    #rospy.loginfo("yaw_deg: %f", yaw_deg)
    yaw = yaw_deg*degrees2rad

    #in AHRS firmware y axis points right, in ROS y axis points left (see REP 103)
    pitch = float(myIMU.raw_pitch)*degrees2rad
    roll  = float(myIMU.raw_roll)*degrees2rad
    #rospy.loginfo("yaw:%f pitch:%f  rool:%f", yaw, pitch, roll)

    # Publish message
    imuMsg.linear_acceleration.x = float(myIMU.raw_ax)*accel_factor
    imuMsg.linear_acceleration.y = float(myIMU.raw_ay)*accel_factor
    imuMsg.linear_acceleration.z = float(myIMU.raw_az)*accel_factor

    imuMsg.angular_velocity.x = float(myIMU.raw_gx)*degrees2rad
    imuMsg.angular_velocity.y = float(myIMU.raw_gy)*degrees2rad
    imuMsg.angular_velocity.z = float(myIMU.raw_gz)*degrees2rad

    #myIMU.get_quatern()
    q = quaternion_from_euler(roll, pitch, yaw)
    imuMsg.orientation.x = q[0]
    imuMsg.orientation.y = q[1]
    imuMsg.orientation.z = q[2]
    imuMsg.orientation.w = q[3]
    imuMsg.header.stamp= rospy.Time.now()
    imuMsg.header.frame_id = link_name
    imuMsg.header.seq = seq
    pub.publish(imuMsg)
    seq = seq + 1
    rate.sleep()