#!/usr/bin/env python
# -*- coding: UTF-8 -*-

# Copyright: 2016-2021 https://www.corvin.cn ROS小课堂
# Author: corvin
# Description: 从IIC接口中读取IMU模块的三轴加速度、三轴角速度、四元数。
# History:
#    20211122: Initial this file.
import smbus
import rospy
import numpy as np

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

    def get_YPRAG(self):
        try:
            rpy_data  = self.i2c.read_i2c_block_data(self.addr, 0x3d, 6)
            axyz_data = self.i2c.read_i2c_block_data(self.addr, 0x34, 6)
            gxyz_data = self.i2c.read_i2c_block_data(self.addr, 0x37, 6)
        except IOError:
            rospy.logerr("Read IMU RPYAG date error !")
        else:
            self.raw_roll  = float(np.short(np.short(rpy_data[1]<<8)|rpy_data[0])/32768.0*180.0)
            self.raw_pitch = float(np.short(np.short(rpy_data[3]<<8)|rpy_data[2])/32768.0*180.0)
            self.raw_yaw   = float(np.short(np.short(rpy_data[5]<<8)|rpy_data[4])/32768.0*180.0)

            self.raw_ax = float(np.short(np.short(axyz_data[1]<<8)|axyz_data[0])/32768.0*16.0)
            self.raw_ay = float(np.short(np.short(axyz_data[3]<<8)|axyz_data[2])/32768.0*16.0)
            self.raw_az = float(np.short(np.short(axyz_data[5]<<8)|axyz_data[4])/32768.0*16.0)

            self.raw_gx = float(np.short(np.short(gxyz_data[1]<<8)|gxyz_data[0])/32768.0*2000.0)
            self.raw_gy = float(np.short(np.short(gxyz_data[3]<<8)|gxyz_data[2])/32768.0*2000.0)
            self.raw_gz = float(np.short(np.short(gxyz_data[5]<<8)|gxyz_data[4])/32768.0*2000.0)

    def get_quatern(self):
        try:
            quat_data = self.i2c.read_i2c_block_data(self.addr, 0x51, 8)
        except IOError:
            rospy.logerr("Read IMU quaternion date error !")
        else:
            self.raw_q0 = float((np.short(np.short(quat_data[1]<<8)|quat_data[0]))/32768.0)
            self.raw_q1 = float((np.short(np.short(quat_data[3]<<8)|quat_data[2]))/32768.0)
            self.raw_q2 = float((np.short(np.short(quat_data[5]<<8)|quat_data[4]))/32768.0)
            self.raw_q3 = float((np.short(np.short(quat_data[7]<<8)|quat_data[6]))/32768.0)
