# qmi8658a.py


from machine import I2C
import time, math, struct

class QMI8658A:
    def __init__(self, i2c: I2C, addr=0x6B):
        self.i2c = i2c
        self.addr = addr
        self._init_device()

        self.roll = 0.0
        self.pitch = 0.0
        self.alpha = 0.90  # 滤波系数
        self.dt = 0.02     # 采样周期

    def _write_reg(self, reg, val):
        self.i2c.writeto_mem(self.addr, reg, bytes([val]))

    def _read_reg(self, reg, n=1):
        return self.i2c.readfrom_mem(self.addr, reg, n)

    def _init_device(self):
        WHO_AM_I = 0x00
        if self._read_reg(WHO_AM_I)[0] != 0x05:
            raise Exception("QMI8658A not detected")

        CTRL1 = 0x02
        CTRL2 = 0x03
        CTRL3 = 0x04
        CTRL7 = 0x08

        self._write_reg(CTRL1, 0x40)                   # 自动地址递增
        self._write_reg(CTRL2, (0b001 << 4) | 0b0110)  # ±4g, ODR 125Hz
        self._write_reg(CTRL3, (0b101 << 4) | 0b0110)  # ±512dps, ODR 125Hz
        self._write_reg(CTRL7, 0x03)                   # 启用 ACC 和 GYRO
        time.sleep(0.1)

    def read_sensor(self):
        AX_L = 0x35
        data = self._read_reg(AX_L, 12)
        ax, ay, az, gx, gy, gz = struct.unpack('<hhhhhh', data)

        ax *= 4 / 32768
        ay *= 4 / 32768
        az *= 4 / 32768
        gx *= 512 / 32768
        gy *= 512 / 32768
        gz *= 512 / 32768

        return ax, ay, az, gx, gy, gz

    def update_orientation(self):
        ax, ay, az, gx, gy, gz = self.read_sensor()

        acc_roll = -math.degrees(math.atan2(ay, az))
        acc_pitch = math.degrees(math.atan2(-ax, math.sqrt(ay*ay + az*az)))

        self.roll += gx * self.dt
        self.pitch += gy * self.dt

        self.roll = self.alpha * self.roll + (1 - self.alpha) * acc_roll
        self.pitch = self.alpha * self.pitch + (1 - self.alpha) * acc_pitch

    def get_orientation(self):
        return self.roll, self.pitch

#     def draw_orientation(self):
#         roll, pitch = self.get_orientation()

#         arrow_up    = "←" if pitch < -10 else " "
#         arrow_down  = "→" if pitch > 10  else " "
#         arrow_left  = "↑" if roll  < -10 else " "
#         arrow_right = "↓" if roll  > 10  else " "

#         print("\033[2J\033[H", end="")  # 清屏
#         print(f"Roll: {roll:+.1f}° | Pitch: {pitch:+.1f}°")
#         print("Orientation:")
#         print(f"   {arrow_left}   ")
#         print(f"{arrow_up} [ ] {arrow_down}")
#         print(f"   {arrow_right}   ")
