# curtain.py
#
# Sample Webots controller file for driving a
# 'curtain' of actuated hanging chains.
#
# No copyright, 2020-2022, Garth Zeglin.  This file is
# explicitly placed in the public domain.

print("loading curtain.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"curtain.py waking up for {name}...")

# Query the number of devices.  The curtain.proto model has one joint actuator
# and one sensor per chain.
num_devices = robot.getNumberOfDevices()
chains = num_devices // 2
print(f"Found {num_devices} devices, assuming {chains} hanging chains.")

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

# Fetch handles for the joint sensors.  The names are generated by the curtain.proto scripting.
joints = [robot.getDevice('joint%d' % (jnum+1)) for jnum in range(chains)]

# Specify the sampling rate for the joint sensors.
for j in joints:
    j.enable(EVENT_LOOP_DT)

# Fetch handles for the position actuator at the top of each chain.
motors = [robot.getDevice('motor%d' % (jnum+1)) for jnum in range(chains)]
for m in motors:
    m.setPosition(0.0)

################################################################
# 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.
    q = [j.getValue() for j in joints]

    # Read any computer keyboard keypresses.  Returns -1 or an integer keycode while a key is held down.
    key = keyboard.getKey()
    if key != -1:
        # convert the integer key number to a lowercase single-character string        
        letter = chr(key).lower()

        # special case: 'p' will enter a passive zero-torque mode
        if letter == 'p':
            for m in motors:
                m.setTorque(0.0)

        # drive to downward reference position
        elif letter == 'd':
            for m in motors:
                m.setPosition(0.0)

        # drive all to front
        elif letter == 'f':
            for m in motors:
                m.setPosition(-0.5)

        # drive all to back
        elif letter == 'b':
            for m in motors:
                m.setPosition(0.5)

        # drive to alternating positions
        elif letter == 'l':
            for i, m in enumerate(motors):
                p = 0.5 if (i&1) == 0 else -0.5
                m.setPosition(p)

        # drive to opposite alternating positions
        elif letter == 'r':
            for i, m in enumerate(motors):
                p = 0.5 if (i&1) == 1 else -0.5
                m.setPosition(p)
                
        # generate a traveling wave (while 'w' is held down)
        elif letter == 'w':
            for i, m in enumerate(motors):
                p = 0.5 * math.sin(1.5 * t + 0.75 * i)
                m.setPosition(p)