import serial
import struct
import time

# ========== CRC16 预计算查找表 ==========
crc_tab = []
for i in range(256):
    crc = 0
    c = i
    for _ in range(8):
        if (crc ^ c) & 1:
            crc = (crc >> 1) ^ 0xA001
        else:
            crc >>= 1
        c >>= 1
    crc_tab.append(crc)


def crc16(data):
    """CRC16-Modbus 校验 (查表法)"""
    crc = 0xFFFF
    for b in data:
        crc = (crc >> 8) ^ crc_tab[(crc ^ b) & 0xFF]
    return crc


# ========== 帧格式常量 (新协议 48 字节) ==========
# 0     2B   帧头 0xAA 0x55
# 2     4B   时间戳 uint32_t (ms, 自启动)
# 6     12B  滤波后 Gyro[3] (float32 × 3)
# 18    12B  滤波后 Accel[3] (float32 × 3)
# 30    16B  四元数 qw, qx, qy, qz (float32 × 4)
# 46    2B   CRC16 (对前 46 字节)
HEADER = b'\xAA\x55'
HEADER_LEN = 2
PAYLOAD_LEN = 44               # uint32 + 3f + 3f + 4f = 4+12+12+16
CRC_LEN = 2
FRAME_LEN = HEADER_LEN + PAYLOAD_LEN + CRC_LEN  # = 48

FIELD_NAMES = ['timestamp_ms',
               'gyro_x', 'gyro_y', 'gyro_z',
               'accel_x', 'accel_y', 'accel_z',
               'qw', 'qx', 'qy', 'qz']


def parse_frame(payload_bytes):
    """解包 44 字节 payload

    Returns:
        dict with keys: timestamp_ms, gyro (3-tuple), accel (3-tuple), quat (4-tuple: qw,qx,qy,qz)
    """
    ts, gx, gy, gz, ax, ay, az, qw, qx, qy, qz = struct.unpack('<I10f', payload_bytes)
    return {
        'timestamp_ms': ts,
        'gyro':  (gx, gy, gz),
        'accel': (ax, ay, az),
        'quat':  (qw, qx, qy, qz),
    }


def quat_to_euler(qw, qx, qy, qz):
    """四元数 → 欧拉角 (ZYX 内旋, deg)

    scipy 格式: [x, y, z, w], STM32 格式: [w, x, y, z]
    """
    from scipy.spatial.transform import Rotation
    r = Rotation.from_quat([qx, qy, qz, qw])
    yaw, pitch, roll = r.as_euler('ZYX', degrees=True)
    return yaw, pitch, roll


def main(port='COM3', baud=115200):
    print(f"打开串口 {port} @ {baud} baud ...")
    ser = serial.Serial(port, baud, timeout=1)
    print("等待 IMU 数据帧 (header: AA 55) ...\n")

    frame_count = 0
    err_count = 0
    last_ts = None

    try:
        while True:
            b = ser.read(1)
            if not b:
                continue
            if b[0] == 0xAA:
                b2 = ser.read(1)
                if not b2 or b2[0] != 0x55:
                    continue

                # 帧头匹配, 读取剩余 46 字节
                frame = ser.read(PAYLOAD_LEN + CRC_LEN)
                if len(frame) < PAYLOAD_LEN + CRC_LEN:
                    print("警告: 帧数据不完整")
                    continue

                payload_bytes = frame[:PAYLOAD_LEN]
                crc_recv = struct.unpack('<H', frame[PAYLOAD_LEN:])[0]

                # CRC 校验 (覆盖 header + payload, 前 46 字节)
                crc_calc = crc16(HEADER + payload_bytes)

                if crc_calc == crc_recv:
                    data = parse_frame(payload_bytes)
                    ts = data['timestamp_ms']
                    gx, gy, gz = data['gyro']
                    ax, ay, az = data['accel']
                    qw, qx, qy, qz = data['quat']
                    yaw, pitch, roll = quat_to_euler(qw, qx, qy, qz)

                    dt_str = ""
                    if last_ts is not None:
                        dt = (ts - last_ts) / 1000.0
                        dt_str = f" dt={dt:.3f}s"
                    last_ts = ts

                    print(f"[#{frame_count:04d}] TS={ts:08d}{dt_str}\n"
                          f"  Gyro:  {gx:8.3f}, {gy:8.3f}, {gz:8.3f}\n"
                          f"  Accel: {ax:7.3f}, {ay:7.3f}, {az:7.3f}\n"
                          f"  Quat:  w={qw:6.3f} x={qx:6.3f} y={qy:6.3f} z={qz:6.3f}\n"
                          f"  Euler: Y{yaw:7.2f} P{pitch:7.2f} R{roll:7.2f}")
                    frame_count += 1
                else:
                    err_count += 1
                    print(f"CRC 校验失败 (期望: 0x{crc_calc:04X}, 接收: 0x{crc_recv:04X}) "
                          f"[累计错误: {err_count}]")

    except KeyboardInterrupt:
        print(f"\n\n停止。共接收 {frame_count} 帧, CRC 错误 {err_count} 次。")
        ser.close()


if __name__ == '__main__':
    import sys
    port = sys.argv[1] if len(sys.argv) > 1 else 'COM3'
    baud = int(sys.argv[2]) if len(sys.argv) > 2 else 115200
    main(port, baud)