import numpy as np
from controller import Robot
from controller import Keyboard

# Instancias
robot = Robot()
keyboard = Keyboard()

# Contantes
timestep = int(robot.getBasicTimeStep())

joints = ['joint_1', 'joint_2', 'joint_3', 'joint_4']
# Initialize motors for each joint
motors = []
for joint in joints:
    motor = robot.getDevice(joint)
    motor.setPosition(0.0)  # Set to infinite position control
    motors.append(motor)

sensors = []
for joint in joints:
    sensor = robot.getDevice(joint + "_sensor")
    sensor.enable(timestep)
    sensors.append(sensor)

def set_joint_positions(positions):
    """Set the positions of the SCARA robot joints."""
    global motors
    for motor, position in zip(motors, positions):
        motor.setPosition(position)

def set_joint_velocity(velocities):
    """Set the positions of the SCARA robot joints."""
    global motors
    for motor, velocity in zip(motors, velocities):
        motor.setVelocity(velocity)

def get_sensor_values():
    """Get the values from the SCARA robot sensors."""
    global sensors
    return [sensor.getValue() for sensor in sensors]

set_joint_velocity([0.3,1,1,1])

position_index = 0
positions = [
    [0,0,-0.097,0],
    [1,1,-0.097,0],
    [0,0,-0.097,0],
    [-1,-1,-0.097,0],
]

keyboard.enable(timestep)
pressedN = False
# Main loop:
# - perform simulation steps until Webots is stopping the controller
while robot.step(timestep) != -1:
    
    key=keyboard.getKey()
    if (key==ord('N')) and not pressedN:
        pressedN = True
    elif (key==-1) and pressedN:
        pressedN = False
        position_index = (position_index+1)%len(positions)
    
    sensor_values = get_sensor_values()
    set_joint_positions(positions[position_index])
 
    pass

# Enter here exit cleanup code.