import smbus
import time
import math
from ahrs.filters import Madgwick

class IMUReader:
    def __init__(self):
        self.bus = smbus.SMBus(1)  # I2C bus
        self.DEVICE_ADDRESS = 0x68  # MPU6050 device address
        self.mpu_init()

        self.filter = Madgwick()
        self.orientation = [0.0, 0.0, 0.0, 1.0]  # Initial quaternion

    def mpu_init(self):
        PWR_MGMT_1 = 0x6B
        SMPLRT_DIV = 0x19
        CONFIG = 0x1A
        GYRO_CONFIG = 0x1B
        ACCEL_CONFIG = 0x1C
        INT_ENABLE = 0x38

        # Write to sample rate register
        self.bus.write_byte_data(self.DEVICE_ADDRESS, SMPLRT_DIV, 7)
        # Write to power management register to wake up the MPU6050
        self.bus.write_byte_data(self.DEVICE_ADDRESS, PWR_MGMT_1, 1)
        # Write to configuration register
        self.bus.write_byte_data(self.DEVICE_ADDRESS, CONFIG, 0)
        # Write to gyroscope configuration register
        self.bus.write_byte_data(self.DEVICE_ADDRESS, GYRO_CONFIG, 24)
        # Write to accelerometer configuration register
        self.bus.write_byte_data(self.DEVICE_ADDRESS, ACCEL_CONFIG, 0)
        # Enable interrupts
        self.bus.write_byte_data(self.DEVICE_ADDRESS, INT_ENABLE, 1)

    def read_raw_data(self, addr):
        # Read two bytes of data from the given address
        high = self.bus.read_byte_data(self.DEVICE_ADDRESS, addr)
        low = self.bus.read_byte_data(self.DEVICE_ADDRESS, addr + 1)
        value = (high << 8) | low
        if value > 32768:
            value = value - 65536
        return value

    def read_data(self):
        ACCEL_XOUT_H = 0x3B
        ACCEL_YOUT_H = 0x3D
        ACCEL_ZOUT_H = 0x3F
        GYRO_XOUT_H = 0x43
        GYRO_YOUT_H = 0x45
        GYRO_ZOUT_H = 0x47

        # Read accelerometer data
        acc_x = self.read_raw_data(ACCEL_XOUT_H)
        acc_y = self.read_raw_data(ACCEL_YOUT_H)
        acc_z = self.read_raw_data(ACCEL_ZOUT_H)

        # Read gyroscope data
        gyro_x = self.read_raw_data(GYRO_XOUT_H)
        gyro_y = self.read_raw_data(GYRO_YOUT_H)
        gyro_z = self.read_raw_data(GYRO_ZOUT_H)

        ax_offset = 0.01467432
        ay_offset = -0.00603101
        az_offset = 0.06171924

        gx_offset = -0.17447328
        gy_offset = 0.08720611
        gz_offset = 0.16423664

        Ax = (acc_x / 16384.0) - ax_offset
        Ay = (acc_y / 16384.0) - ay_offset
        Az = (acc_z / 16384.0) - az_offset
        Ax = 9.81 * Ax
        Ay = 9.81 * Ay
        Az = 9.81 * Az

        Gx = (gyro_x / 131.0) - gx_offset
        Gy = (gyro_y / 131.0) - gy_offset
        Gz = (gyro_z / 131.0) - gz_offset   

        Gx = Gx / 180 * math.pi
        Gy = Gy / 180 * math.pi
        Gz = Gz / 180 * math.pi

        self.orientation = self.filter.updateIMU(self.orientation,
                                                 [Gx, Gy, Gz],
                                                 [Ax / 9.81, Ay / 9.81, Az / 9.81])
        return Ax, Ay, Az, Gx, Gy, Gz, self.orientation

    def start(self):
        print("Starting IMU data collection...")
        while True:
            Ax, Ay, Az, Gx, Gy, Gz, orientation = self.read_data()

            # Debugging: Raw data
            print(f"Raw Accelerometer: Ax={Ax:.2f}, Ay={Ay:.2f}, Az={Az:.2f}")
            print(f"Raw Gyroscope: Gx={Gx:.2f}, Gy={Gy:.2f}, Gz={Gz:.2f}")

            # Debugging: Orientation
            print(f"Orientation (quaternion): x={orientation[0]:.2f}, y={orientation[1]:.2f}, "
                  f"z={orientation[2]:.2f}, w={orientation[3]:.2f}")

            # Test stability
            print("Orientation (Euler angles):")
            roll = math.atan2(2.0 * (orientation[3] * orientation[0] + orientation[1] * orientation[2]),
                              1.0 - 2.0 * (orientation[0]**2 + orientation[1]**2))
            pitch = math.asin(2.0 * (orientation[3] * orientation[1] - orientation[2] * orientation[0]))
            yaw = math.atan2(2.0 * (orientation[3] * orientation[2] + orientation[0] * orientation[1]),
                             1.0 - 2.0 * (orientation[1]**2 + orientation[2]**2))
            print(f"Roll={roll:.2f}, Pitch={pitch:.2f}, Yaw={yaw:.2f}")

            time.sleep(0.5)


if __name__ == "__main__":
    imu_reader = IMUReader()
    imu_reader.start()