# dbl_pend.py
#
# Sample Webots controller file for driving a
# fully actuated double pendulum.  With both actuators
# set to zero torque the device will act like a chaotic
# double pendulum. Otherwise this code implements
# PD position control to drive towards a pose.
#
# No copyright, 2020-2022, Garth Zeglin.  This file is
# explicitly placed in the public domain.

print("loading dbl_pend.py...")

# Import the Webots simulator API.
from controller import Robot
from controller import Keyboard

# Import the standard Python math library.
import math

# Define the time step in milliseconds between
# controller updates.
EVENT_LOOP_DT = 20

################################################################

# Request a proxy object representing the robot to control.
robot = Robot()
name = robot.getName()
print(f"dbl_pend.py waking up for {name}...")

# Enable computer keyboard input for user control.
keyboard = Keyboard()
keyboard.enable(EVENT_LOOP_DT)

# Fetch handles for the joint sensors.
j1 = robot.getDevice('joint1')
j2 = robot.getDevice('joint2')

# Specify the sampling rate for the joint sensors.
j1.enable(EVENT_LOOP_DT)
j2.enable(EVENT_LOOP_DT)

# Fetch handle for the 'base' joint motor.  In this
# example the motor will be controlled as a torque
# motor, bypassing the Webots PID control.
motor1 = robot.getDevice('motor1')
motor1.setTorque(0.0)

motor2 = robot.getDevice('motor2')
motor2.setTorque(0.0)

# set properties in RotationalMotor nodes:
motor1.setAvailableTorque(4)
motor2.setAvailableTorque(1)

# Specify proportional and derivative (damping) gains
# for the position controllers in the main event loop.
P_gain1 = 8.0     # units are N-m / radian
D_gain1 = 0.2     # units are N-m / (radian/sec)

P_gain2 = 4.0     # units are N-m / radian
D_gain2 = 0.1     # units are N-m / (radian/sec)

# Soft torque limits to avoid error messages.
max_tau1 = 4.0  # units are N-m
max_tau2 = 1.0  # units are N-m

# Initial position targets.
q_d1 = None
q_d2 = None

# The controller estimates velocities using finite
# differencing of the position sensors; these variables
# hold the previous state.
last_q1 = 0
last_q2 = 0

################################################################
# Define reference poses mapped to keyboard keys.

poses = {
    't' : (0.0,  0.0),          # straight up
    'r' : (0.5, -1.0),          # elbow to machine right (-Y)
    'l' : (-0.5, 1.0),          # elbow to machine left (+Y)
    'b' : (math.pi, -math.pi),  # elbow at bottom, link2 pointing up
    'w' : (2*math.pi, 0.0),     # straight up, but wound a full positive rotation
}

# Set joint position targets to a reference pose indicated by a single character string.
def set_pose(key):
    global q_d1, q_d2
    # special case: 'p' will enter a passive zero-torque mode
    if key == 'p':
        q_d1 = None
        q_d2 = None
    else:
        # look for a reference pose
        pose = poses.get(key)
        if pose is not None:
            q_d1 = pose[0]
            q_d2 = pose[1]

################################################################
# Metronomic choreography for automatic machine.
def set_automatic_targets(t):
    phase = math.fmod(t, 5.0) / 5.0
    if phase < 0.5:
        set_pose('l')
    else:
        set_pose('r')

################################################################

# Run an event loop until the simulation quits,
# indicated by the step function returning -1.
while robot.step(EVENT_LOOP_DT) != -1:

    # Read simulator clock time.
    t = robot.getTime()

    # Read the new joint positions.
    q1 = j1.getValue()
    q2 = j2.getValue()

    # Estimate current velocities in radians/sec using finite differences.
    qd1 = (q1 - last_q1) / (0.001 * EVENT_LOOP_DT)
    qd2 = (q2 - last_q2) / (0.001 * EVENT_LOOP_DT)
    last_q1 = q1
    last_q2 = q2

    # Read any computer keyboard keypresses.  Returns -1 or an integer keycode while a key is held down.
    if name == 'left':
        key = keyboard.getKey()
        if key != -1:
            # convert the integer key number to a lowercase single-character string,
            # then set the target pose indicated by the key
            set_pose(chr(key).lower())

    else:
        # the other devices run an automatic timed sequence
        set_automatic_targets(t)

    # Calculate new motor torques, limit them, and apply them to the system.
    if q_d1 is None:
        tau1 = 0.0
    else:
        tau1 = P_gain1 * (q_d1 - q1) - D_gain1 * qd1
        tau1 = min(max(tau1, -max_tau1), max_tau1)

    if q_d2 is None:
        tau2 = 0.0
    else:
        tau2 = P_gain2 * (q_d2 - q2) - D_gain2 * qd2
        tau2 = min(max(tau2, -max_tau2), max_tau2)

    motor1.setTorque(tau1)
    motor2.setTorque(tau2)