from machine import Pin, PWM
import utime

# Define the servo motor pins
servo1 = PWM(Pin(15)) #(Base)
servo2 = PWM(Pin(19)) #(Arm)
servo3 = PWM(Pin(18)) #(Wrist)
servo4 = PWM(Pin(17)) #(Hand)

# Set the frequency
servo1.freq(100) #(Base)
servo2.freq(100) #(Arm)
servo3.freq(100) #(Wrist)
servo4.freq(100) #(Hand)

# Set initial servo angles
angle1 = 0
angle2 = 0
angle3 = 0
angle4 = 0

# Function to move the servo motor
def move_servo(servo, angle):
    duty = int((angle / 180) * (2500 - 500) + 500)
    servo.duty_u16(duty * 65535 // 3300)

# Function to control the robot arm using keyboard
def control_arm():
    global angle1, angle2, angle3, angle4
    keys = ['w', 's', 'a', 'd', 'g', 'o', 'q', 'e']
    key_map = {
        'w': (servo1, 10),
        's': (servo1, -10),
        'a': (servo2, 10),
        'd': (servo2, -10),
        'g': (servo3, 10),
        'o': (servo3, -10),
        'q': (servo4, 10),
        'e': (servo4, -10)
    }
    for key in keys:
        Pin(key, Pin.IN, Pin.PULL_UP) # Set the pin as input with pull-up resistor
    while True:
        for key in keys:
            if Pin(key).value() == 0:
                servo, change = key_map[key]
                if key == 'w' and angle1 == 180:
                    continue
                elif key == 's' and angle1 == 0:
                    continue
                elif key == 'a' and angle2 == 180:
                    continue
                elif key == 'd' and angle2 == 0:
                    continue
                elif key == 'g' and angle3 == 180:
                    continue
                elif key == 'o' and angle3 == 0:
                    continue
                elif key == 'q' and angle4 == 180:
                    continue
                elif key == 'e' and angle4 == 0:
                    continue
                angle = min(180, max(0, eval(f'angle{key_map[key][0].pwm_pin()} + change')))
                eval(f'move_servo(servo{key_map[key][0].pwm_pin()}, {angle})')
                print(f"Servo {key_map[key][0].pwm_pin()} moved")
                eval(f'angle{key_map[key][0].pwm_pin()} = angle')
                utime.sleep(0.1)

# Start controlling the robot arm
control_arm()