import pygame
import serial
import time

# Constants
COM_PORT = 'COM6'  # Replace with your Arduino's serial port
BAUD_RATE = 9600

#JOYSTICKs IDs
# 0: left joystick, horizontal
# 1: left joystick, vertical
# 2: right joystick, horizontal
# 3: right joystick, vertical

BASE_ROTATION_AXIS_ID = 0
ARM_ROTATION_AXIS_ID = 1

GRIPPER_ROLL_AXIS_ID = 2
GRIPPER_PITCH_AXIS_ID = 3

#BUTTONs IDs
# 0: "1" button
# 1: "2" button
# 2: "3" button
# 3: "4" button
# 4:
# 5:
# 6: left "big" trigger
# 7: right "big" trigger
# 8:
# 9:
# 10:
# 11:

ELBOW_DOWN_BUTTON = 1  # Button ID for elbow-down
ELBOW_UP_BUTTON = 3  # Button ID for elbow-up

GRIPPER_CLOSE_BUTTON = 6  # Button ID for gripper close
GRIPPER_OPEN_BUTTON = 7  # Button ID for gripper open

#NUMHATS can also be used (left pad)


# Serial setup
try:
    ser = serial.Serial(COM_PORT, BAUD_RATE)
    time.sleep(2)  # Allow time for the connection to establish
    print(f"Connected to {COM_PORT}")
except serial.SerialException as e:
    print(f"Error: {e}")
    ser = None

# Initialize Pygame
pygame.init()
pygame.joystick.init()

# Check for joystick (controller)
if pygame.joystick.get_count() == 0:
    print("No joystick detected!")
    exit()

joystick = pygame.joystick.Joystick(0)
joystick.init()

print(f"Using joystick: {joystick.get_name()}")
#print(f"number of buttons: {joystick.get_numbuttons()}");

## Map joystick values to ranges
# def map_value(value, from_min, from_max, to_min, to_max):
    # return to_min + (value - from_min) * (to_max - to_min) / (from_max - from_min)


# Main loop
running = True

# q1: base rotation
# q2: arm rotation
# q6: elbow rotation
q1_steps, q2_steps, q6_steps, gripper_roll, gripper_pitch, gripper_closing_steps = 0, 0, 0, 0, 0, 0

while running:
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False

    # Read joystick axes
    # Scale by 5 to get faster movement
    
    q1_steps += -5*joystick.get_axis(BASE_ROTATION_AXIS_ID)
    q2_steps += -5*joystick.get_axis(ARM_ROTATION_AXIS_ID)
    gripper_roll += -5*joystick.get_axis(GRIPPER_ROLL_AXIS_ID)
    gripper_pitch += -5*joystick.get_axis(GRIPPER_PITCH_AXIS_ID)
    
    # Read button states
    # +-5 for faster movement
    if joystick.get_button(ELBOW_DOWN_BUTTON):
        q6_steps-=5
    elif joystick.get_button(ELBOW_UP_BUTTON):
        q6_steps+=5
    if joystick.get_button(GRIPPER_CLOSE_BUTTON):
        gripper_closing_steps-=5
    elif joystick.get_button(GRIPPER_OPEN_BUTTON):
        gripper_closing_steps+=5
    

    # This might need to be changed as wrong calibration will prevent the robot from moving to its full range
    # moreover, the grip cannot be calibrated as it sometimes "jumps" when the string does a full rotation around the motor
    
    # Clamp q6_steps values to valid range
    q1_steps = max(0, min(2800, q1_steps))
    # Clamp q6_steps values to valid range
    q2_steps = max(0, min(2000, q2_steps))
    # Clamp gripper values to valid range
##    gripper_closing_steps = max(0, gripper_closing_steps)


    # Send values to Arduino
    values = [int(q1_steps), int(q2_steps), int(gripper_closing_steps), int(gripper_roll), int(gripper_pitch), int(q6_steps)]
    values_str = ','.join(map(str, values))
    if ser:
        try:
            ser.write((values_str + '\n').encode())
            print(f"Sent: {values_str}")
            if ser.in_waiting >0:
                line = ser.readline().decode("utf-8").strip()
                print(f"<- {line}")
        except serial.SerialException as e:
            print(f"Serial write error: {e}")
            
    # Limit update rate
    pygame.time.wait(50)

# Clean up
if ser:
    ser.close()
pygame.quit()