#!/usr/bin/env python3
"""
imu_test.py - raw MPU-6500 / MPU-9250 readout (GY-6500, GY-9250 boards).
No filtering, no calibration, no sign flips.
Just shows what each axis does so we can map them to the robot.
Ctrl-C to quit.
"""
import time
from smbus2 import SMBus

BUS = 1

# MPU-6500 / MPU-9250 registers
CONFIG        = 0x1A   # gyro low-pass filter
GYRO_CONFIG   = 0x1B   # gyro range
ACCEL_CONFIG  = 0x1C   # accel range
ACCEL_CONFIG2 = 0x1D   # accel low-pass filter
ACCEL_XOUT_H  = 0x3B   # accel x/y/z, temp, gyro x/y/z = 14 bytes
PWR_MGMT_1    = 0x6B
WHO_AM_I      = 0x75

KNOWN = {0x70: "MPU-6500", 0x71: "MPU-9250", 0x73: "MPU-9255", 0x68: "MPU-6050"}


def s16(h, l):
    v = (h << 8) | l
    return v - 65536 if v & 0x8000 else v


bus = SMBus(BUS)

addr = None
for a in (0x68, 0x69):
    try:
        who = bus.read_byte_data(a, WHO_AM_I)
        name = KNOWN.get(who, "unknown")
        print(f"0x{a:02X}: WHO_AM_I = 0x{who:02X}  ({name})")
        if who in KNOWN and addr is None:
            addr = a
    except OSError:
        print(f"0x{a:02X}: no response")

if addr is None:
    raise SystemExit("No MPU found. Check wiring / run i2cdetect -y 1")

w = lambda reg, val: bus.write_byte_data(addr, reg, val)
w(PWR_MGMT_1, 0x80); time.sleep(0.1)   # reset
w(PWR_MGMT_1, 0x01); time.sleep(0.05)  # wake up, PLL clock
w(CONFIG, 0x03)                        # gyro low-pass ~41 Hz
w(GYRO_CONFIG, 0x08)                   # +-500 deg/s
w(ACCEL_CONFIG, 0x00)                  # +-2 g
w(ACCEL_CONFIG2, 0x03)                 # accel low-pass ~41 Hz
time.sleep(0.05)

print("\n   ax      ay      az    |a|   |     gx       gy       gz   (g, deg/s)")
try:
    while True:
        d = bus.read_i2c_block_data(addr, ACCEL_XOUT_H, 14)
        ax, ay, az = (s16(d[i], d[i + 1]) / 16384.0 for i in (0, 2, 4))
        gx, gy, gz = (s16(d[i], d[i + 1]) / 65.5 for i in (8, 10, 12))   # skip temp at 6-7
        mag = (ax * ax + ay * ay + az * az) ** 0.5
        print(f"{ax:+6.2f}  {ay:+6.2f}  {az:+6.2f}  {mag:4.2f}  |  {gx:+7.1f}  {gy:+7.1f}  {gz:+7.1f}")
        time.sleep(0.2)
except KeyboardInterrupt:
    print()
finally:
    bus.close()