import time
import board
import busio
import usb_hid
import struct
import digitalio
import analogio
import adafruit_mpu6050
import math

# https://forums.adafruit.com/viewtopic.php?t=194490
if True:        # optionally wiggle SCL to clear I2C slave devices
  print("***** Clear I2C bus *****")
  scl = digitalio.DigitalInOut(board.GP1)
  scl.direction = digitalio.Direction.OUTPUT  # want open drain scl.drive_mode = digitalio.DriveMode.OPEN_DRAIN
  for n in range(1,20):
    scl.value = 0
    time.sleep(0.0001)
    scl.value = 1
    time.sleep(0.0001)
  scl.direction = digitalio.Direction.INPUT
  scl.deinit()

# Initialize I2C and MPU6050
i2c = busio.I2C(board.GP1, board.GP0)
mpu = adafruit_mpu6050.MPU6050(i2c)

# Configure MPU6050
mpu.gyro_range = adafruit_mpu6050.GyroRange.RANGE_2000_DPS
mpu.filter_bandwidth = adafruit_mpu6050.Bandwidth.BAND_44_HZ
mpu.sample_rate_divisor = 9  # Set sample rate to 100Hz

# Create a custom HID device
gamepad = usb_hid.devices[0]  # Assuming the first device is the gamepad

# Button setup (digital inputs)
buttons = {
    1: digitalio.DigitalInOut(board.GP2),
    2: digitalio.DigitalInOut(board.GP3),
    3: digitalio.DigitalInOut(board.GP4),
    4: digitalio.DigitalInOut(board.GP5),
    5: digitalio.DigitalInOut(board.GP6),
    6: digitalio.DigitalInOut(board.GP7),
    7: digitalio.DigitalInOut(board.GP8),
}

for button in buttons.values():
    button.direction = digitalio.Direction.INPUT
    button.pull = digitalio.Pull.UP

# Joystick setup (analog inputs)
joy_x_axis = analogio.AnalogIn(board.GP26)
joy_y_axis = analogio.AnalogIn(board.GP27)

# Calibration (replace with actual calibration values)
gyro_x_offset = 0.0
gyro_y_offset = 0.0
gyro_z_offset = 0.0

# Complementary filter parameters
ALPHA = 0.98  # Weight for gyroscope data (0 < ALPHA < 1)
DT = 0.005  # Sampling time (200Hz)

# Initialize angles
pitch = 0.0
roll = 0.0
yaw = 0.0

# Function to calculate angles from accelerometer data
def calculate_angles(accel_x, accel_y, accel_z):
    pitch = math.atan2(accel_y, math.sqrt(accel_x**2 + accel_z**2)) * (180.0 / math.pi)
    roll = math.atan2(-accel_x, accel_z) * (180.0 / math.pi)
    return pitch, roll

# Function to apply complementary filter
def complementary_filter(gyro_data, accel_data, prev_angles, alpha, dt):
    # Calculate angles from accelerometer
    accel_pitch, accel_roll = calculate_angles(*accel_data)
    
    # Integrate gyroscope data to get angles
    gyro_pitch = prev_angles[0] + gyro_data[0] * dt
    gyro_roll = prev_angles[1] + gyro_data[1] * dt
    gyro_yaw = prev_angles[2] + gyro_data[2] * dt  # Integrate Z-axis for yaw
    
    # Combine gyroscope and accelerometer angles using complementary filter
    pitch = alpha * gyro_pitch + (1 - alpha) * accel_pitch
    roll = alpha * gyro_roll + (1 - alpha) * accel_roll
    yaw = gyro_yaw  # Yaw is purely from gyroscope
    
    return pitch, roll, yaw

def read_buttons():
    """Reads button states and updates the bitmask for HID report."""
    buttons_state = 0
    for button_num, button in buttons.items():
        if not button.value:  # Button pressed (active low)
            buttons_state |= 1 << (button_num - 1)
    return buttons_state

def read_joysticks():
    """Reads analog joystick values and scales them to -127 to 127."""
    joy_x = max(-127, min(127, ((joy_x_axis.value - 32768) * 127) // 32768))
    joy_y = max(-127, min(127, ((joy_y_axis.value - 32768) * 127) // 32768))
    return joy_x, joy_y

print("MPU6050 Game Controller Running!")

while True:
    while not i2c.try_lock():
        pass
    
    print("I2C addresses found:", [hex(x) for x in i2c.scan()])
    i2c.unlock()
    
    try:
        gyro_x, gyro_y, gyro_z = mpu.gyro
        accel_x, accel_y, accel_z = mpu.acceleration
    except OSError as e:
        print(f"Sensor read failed: {e}")
        time.sleep(0.1)
        
        continue
    
    # Read and calibrate gyro data
    # gyro_x, gyro_y, gyro_z = (0, 0, 0)
    # gyro_x, gyro_y, gyro_z = mpu.gyro
    gyro_x -= gyro_x_offset
    gyro_y -= gyro_y_offset
    gyro_z -= gyro_z_offset
    # accel_x, accel_y, accel_z = (0, 0, 0) 
    # accel_x, accel_y, accel_z = mpu.acceleration

    # Apply complementary filter to calculate pitch, roll, and yaw
    pitch, roll, yaw = complementary_filter(
        (gyro_x, gyro_y, gyro_z),  # Gyroscope data (angular velocity in rad/s)
        (accel_x, accel_y, accel_z),  # Accelerometer data
        (pitch, roll, yaw),  # Previous angles
        ALPHA,  # Filter parameter
        DT  # Sampling time
    )

    # Scale angles to -127 to 127
    pitch_scaled = max(-127, min(127, int(pitch)))
    roll_scaled = max(-127, min(127, int(roll)))
    yaw_scaled = max(-127, min(127, int(yaw * 100)))

    # Read button and joystick states
    buttons_state = read_buttons()
    joy_x, joy_y = read_joysticks()

    # Send HID report
    report = bytearray(8)
    struct.pack_into("<H", report, 0, buttons_state)  # Buttons (2 bytes)
    struct.pack_into("<bb", report, 2, joy_x, joy_y)  # Joystick X, Y (X and Y axes)
    struct.pack_into("<bbb", report, 4, pitch_scaled, roll_scaled, yaw_scaled)  # Pitch, Roll, Yaw (scaled to -127 to 127)

    gamepad.send_report(report)

    time.sleep(DT)