import smbus
import time

class BNO055:
    def __init__(self, i2c_bus=1, address=0x28):
        self.bus = smbus.SMBus(i2c_bus)
        self.address = address
        self.init_bno055()

    def write_register(self, register, value):
        self.bus.write_byte_data(self.address, register, value)

    def read_register(self, register, length=1):
        if length == 1:
            return self.bus.read_byte_data(self.address, register)
        else:
            return self.bus.read_i2c_block_data(self.address, register, length)

    def init_bno055(self):
        # Switch to CONFIG mode to configure sensor
        self.write_register(0x3D, 0x00)
        time.sleep(0.05)

        # Set power mode to Normal
        self.write_register(0x3E, 0x00)

        # Set to NDOF mode (Full sensor fusion for orientation)
        self.write_register(0x3D, 0x0C)
        time.sleep(0.05)

    def read_euler_angles(self):
        # Euler angles (yaw, pitch, roll) registers
        data = self.read_register(0x1A, 6)
        yaw = (data[1] << 8 | data[0]) / 16.0
        pitch = (data[3] << 8 | data[2]) / 16.0
        roll = (data[5] << 8 | data[4]) / 16.0
        return yaw, pitch, roll

if __name__ == "__main__":
    sensor = BNO055()
    print("Starting BNO055 data collection...")
    while True:
        yaw, pitch, roll = sensor.read_euler_angles()
        print(f"Yaw: {yaw:.2f}, Pitch: {pitch:.2f}, Roll: {roll:.2f}")
        time.sleep(0.5)