import time
import board
import busio
import usb_cdc
import digitalio
import analogio
import adafruit_mpu6050
import math

# Setup I2C for MPU6050
i2c = busio.I2C(board.GP1, board.GP0)
mpu = adafruit_mpu6050.MPU6050(i2c)

# MPU6050 config
mpu.gyro_range = adafruit_mpu6050.GyroRange.RANGE_2000_DPS
mpu.filter_bandwidth = adafruit_mpu6050.Bandwidth.BAND_44_HZ
mpu.sample_rate_divisor = 9  # ~100Hz

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

# Buttons
button_pins = [board.GP2, board.GP3, board.GP4, board.GP5, board.GP6, board.GP7, board.GP8, board.GP9, board.GP10]
buttons = []

for pin in button_pins:
    btn = digitalio.DigitalInOut(pin)
    btn.direction = digitalio.Direction.INPUT
    btn.pull = digitalio.Pull.UP
    buttons.append(btn)

# Complementary filter vars
pitch, roll, yaw = 0.0, 0.0, 0.0
ALPHA = 0.98
DT = 0.01  # 100Hz

def calculate_angles(ax, ay, az):
    pitch = math.atan2(ay, math.sqrt(ax**2 + az**2)) * 180.0 / math.pi
    roll = math.atan2(-ax, az) * 180.0 / math.pi
    return pitch, roll

def complementary_filter(gyro, accel, prev_angles):
    accel_pitch, accel_roll = calculate_angles(*accel)
    gyro_pitch = prev_angles[0] + gyro[0] * DT
    gyro_roll = prev_angles[1] + gyro[1] * DT
    gyro_yaw = prev_angles[2] + gyro[2] * DT

    pitch = ALPHA * gyro_pitch + (1 - ALPHA) * accel_pitch
    roll = ALPHA * gyro_roll + (1 - ALPHA) * accel_roll
    yaw = gyro_yaw

    return pitch, roll, yaw

def read_joysticks():
    jx = ((joy_x_axis.value - 32768) * 127) // 32768
    jy = ((joy_y_axis.value - 32768) * 127) // 32768
    return max(-127, min(127, jx)), max(-127, min(127, jy))

def read_buttons():
    return [int(not b.value) for b in buttons]  # 1 = pressed

print("Streaming controller data over USB CDC...")

while True:
    try:
        gx, gy, gz = mpu.gyro
        ax, ay, az = mpu.acceleration
    except Exception as e:
        print(f"Sensor error: {e}")
        continue

    pitch, roll, yaw = complementary_filter((gx, gy, gz), (ax, ay, az), (pitch, roll, yaw))

    # Scale angles
    pitch_scaled = max(-127, min(127, int(pitch)))
    roll_scaled = max(-127, min(127, int(roll)))
    yaw_scaled = max(-127, min(127, int(yaw * 10)))  # Scale for readability

    jx, jy = read_joysticks()
    button_states = read_buttons()

    # Format: x,y,pitch,roll,yaw,b1,b2,...,b9
    msg = f"{jx},{jy},{pitch_scaled},{roll_scaled},{yaw_scaled}," + ",".join(map(str, button_states)) + "\n"
    usb_cdc.data.write(msg.encode('utf-8'))

    time.sleep(DT)