import ikpy.chain as ik
import numpy as np
from controller import Supervisor
from controller import Keyboard

# Instancias
robot = Supervisor()
keyboard = Keyboard()
tcp = robot.getFromDef("TCP")

# Contantes
TIME_STEP = int(robot.getBasicTimeStep())
IKPY_MAX_ITERATIONS = 4

joints = ['joint_1', 'joint_2', 'joint_3', 'joint_4']
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(TIME_STEP)
    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([1,1,1,1])

chain = ik.Chain.from_urdf_file("../../urdf/Epson_SSR-H803N.urdf", active_links_mask=[
                                False, True, True, True, True, False, False], last_link_vector=[0, 0, 0])

x_goal = 0.78
y_goal = 0.0
z_goal = 0.3

keyboard.enable(TIME_STEP)
pressedN = False
# Main loop:
# - perform simulation steps until Webots is stopping the controller
while robot.step(TIME_STEP) != -1:
    
    key=keyboard.getKey()
    if(key==49):
        x_goal = x_goal + 0.001
    elif(key==81):
        x_goal = x_goal - 0.001
    elif(key==50):
        y_goal = y_goal + 0.001
    elif(key==87):
        y_goal = y_goal - 0.001
    elif(key==51):
        z_goal = z_goal + 0.001
    elif(key==69):
        z_goal = z_goal - 0.001
    
    #print(key)
    
    initial_position = [0] + get_sensor_values() + [0,0]
    ik_solution = chain.inverse_kinematics(
        [-y_goal, x_goal, z_goal],
        [0,0,1],
        orientation_mode="Z",
        initial_position=initial_position,
        max_iter=IKPY_MAX_ITERATIONS)[1:5]
    
    
    sensor_values = get_sensor_values()
    set_joint_positions(ik_solution)
    tcp_pos = tcp.getPosition()
    print("[qi]:",np.round(ik_solution,decimals=3),"\t[x,y,z]:",np.round(tcp_pos,decimals=3))
    pass

# Enter here exit cleanup code.