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

# Copyright: 2016-2022 https://www.corvin.cn ROS小课堂
# Description: 从串口中读取6DOF IMU模块的三轴加速度、角度、四元数,
#   然后将其打印输出.这里需要注意IMU模块使用的坐标系为ENU东北天坐标系.
#   imu模块的串口波特率是固定的57600bps,对于树莓派的串口这里使用的是ttyAMA0,
#   就是排针处的硬件串口,数据通信更为稳定.要想使用该串口需要交换树莓派默认
#   的ttyS0和ttyAMA0，使ttyAMA0为主串口,这里需要先在rasp-config中关闭串口
#   登录功能,启用串口才可以.
# Author: corvin
# History:
#    20220525:Initial this file.

import serial 


ACCData    = [0.0]*8
GYROData   = [0.0]*8
AngleData  = [0.0]*8

FrameState = 0  #保存当前数据帧类型
Bytenum    = 0  #保存读取到数据帧的第几个字节
CheckSum   = 0  #数据帧求和校验位

acc_data     = [0.0]*3
angular_data = [0.0]*3
rpy_data     = [0.0]*3

# 使用树莓派串口与IMU板进行通信
serial_port="/dev/ttyAMA0"


#对数据进行划分,各自读到对应的数组里
def DueData(inputdata):
    #在局部修改全局变量，要进行global的定义
    global  FrameState 
    global  Bytenum
    global  CheckSum
    global  acc_data
    global  angular_data 
    global  rpy_data

    for data in inputdata:  #在输入的数据进行遍历
        if FrameState == 0:  #未读取到数据帧头
            if data == 0x55 and Bytenum == 0: #读取到数据帧头,增大bytenum
                CheckSum = data
                Bytenum = 1
                continue
            elif data == 0x51 and Bytenum == 1: #在byte不为0 且识别到 0x51的时候,改变frame
                CheckSum += data
                FrameState = 1
                Bytenum = 2
            elif data == 0x52 and Bytenum == 1:
                CheckSum += data
                FrameState = 2
                Bytenum = 2
            elif data==0x53 and Bytenum==1:
                CheckSum+=data
                FrameState=3
                Bytenum=2
        elif FrameState == 1: # acc    #已确定数据代表加速度
            if Bytenum<10:            # 读取8个数据
                ACCData[Bytenum-2]=data # 从0开始
                CheckSum+=data
                Bytenum+=1
            else:
                if data == (CheckSum&0xff):  #假如校验位正确
                    a = get_acc(ACCData)
                    print("Acc:", a)
                CheckSum = 0         #各数据归零，进行新的循环判断
                Bytenum = 0
                FrameState = 0
        elif FrameState==2: # gyro
            if Bytenum<10:
                GYROData[Bytenum-2]=data
                CheckSum+=data
                Bytenum+=1
            else:
                if data == (CheckSum&0xff):
                    w = get_gyro(GYROData)
                    print("Angular:", w)
                CheckSum=0
                Bytenum=0
                FrameState=0
        elif FrameState==3: # angle
            if Bytenum<10:
                AngleData[Bytenum-2]=data
                CheckSum+=data
                Bytenum+=1
            else:
                if data == (CheckSum&0xff):
                    Angle = get_angle(AngleData)
                    print("Angle(deg):", Angle)
                CheckSum=0
                Bytenum=0
                FrameState=0


def get_acc(datahex):
    axl = datahex[0]
    axh = datahex[1]
    ayl = datahex[2]
    ayh = datahex[3]
    azl = datahex[4]
    azh = datahex[5]
    
    k_acc = 16.0
    acc_x = (axh << 8 | axl) / 32768.0 * k_acc
    acc_y = (ayh << 8 | ayl) / 32768.0 * k_acc
    acc_z = (azh << 8 | azl) / 32768.0 * k_acc

    if acc_x >= k_acc:
        acc_x -= 2 * k_acc

    if acc_y >= k_acc:
        acc_y -= 2 * k_acc

    if acc_z >= k_acc:
        acc_z-= 2 * k_acc

    return acc_x,acc_y,acc_z


def get_gyro(datahex):
    wxl = datahex[0]
    wxh = datahex[1]
    wyl = datahex[2]
    wyh = datahex[3]
    wzl = datahex[4]
    wzh = datahex[5]

    k_gyro = 2000.0
    gyro_x = (wxh << 8 | wxl) / 32768.0 * k_gyro
    gyro_y = (wyh << 8 | wyl) / 32768.0 * k_gyro
    gyro_z = (wzh << 8 | wzl) / 32768.0 * k_gyro

    if gyro_x >= k_gyro:
        gyro_x -= 2 * k_gyro

    if gyro_y >= k_gyro:
        gyro_y -= 2 * k_gyro

    if gyro_z >=k_gyro:
        gyro_z-= 2 * k_gyro

    return gyro_x,gyro_y,gyro_z


def get_angle(datahex):
    rxl = datahex[0]
    rxh = datahex[1]
    ryl = datahex[2]
    ryh = datahex[3]
    rzl = datahex[4]
    rzh = datahex[5]

    k_angle = 180.0
    angle_x = (rxh << 8 | rxl) / 32768.0 * k_angle
    angle_y = (ryh << 8 | ryl) / 32768.0 * k_angle
    angle_z = (rzh << 8 | rzl) / 32768.0 * k_angle

    if angle_x >= k_angle:
        angle_x -= 2 * k_angle

    if angle_y >= k_angle:
        angle_y -= 2 * k_angle

    if angle_z >=k_angle:
        angle_z-= 2 * k_angle

    return angle_x,angle_y,angle_z


if __name__=='__main__':
    # 开启树莓派串口开始获取数据,这里波特率是固定的57600bps
    mySerial = serial.Serial(serial_port, 57600, timeout=0.5)  
    print("Now IMU module is working...")

    while True:
        # 开始从串口读取44字节数据,共计4帧有效数据进行处理
        datahex = mySerial.read(44)
        DueData(datahex)