# wobbly_thd.py

# Sample Webots controller file for driving the wobbly diff-drive mobile robot.
# This sample uses a separate thread to specify sequential performance actions.

# No copyright, 2020-2026, Garth Zeglin.  This file is
# explicitly placed in the public domain.

print("loading wobbly_thd.py...")

# Import the Webots simulator API.
from controller import Robot

# Import standard Python libraries.
import math, random, time, threading

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

################################################################
# Define a background thread for sequential performance programming.  This runs
# asynchronously and may block using sleep.  It communicates with the real-time
# robot event loop through specific thread-safe methods.  It is set to daemon
# status so it will automatically exit once the main event thread finishes.

class Script(threading.Thread):
    def __init__(self, robot):
        super(Script, self).__init__()
        self.robot = robot
        self.daemon = True

    def run(self):
        print("script thread starting.")
        time.sleep(2)

        while True:
            print("starting script loop iteration.")
            self.robot.go_forward(0.2)
            time.sleep(2)

            self.robot.go_rotate(0.5)
            time.sleep(2)

            self.robot.go_forward(-0.2)
            time.sleep(2)

            self.robot.go_rotate(-0.5)
            time.sleep(2)

            self.robot.go_forward(0.0)
            time.sleep(4)

################################################################
class Wobbly(Robot):
    def __init__(self):

        super(Wobbly, self).__init__()
        self.robot_name = self.getName()
        print("%s: controller connected." % (self.robot_name))

        # Attempt to randomize the random library sequence.
        random.seed(time.time())

        # Initialize geometric constants.  These should match
        # the current geometry of the robot.
        self.wheel_radius = 0.1
        self.axle_length  = 0.14

        # Fetch handles for the wheel motors
        self.l_motor = self.getDevice('left wheel motor')
        self.r_motor = self.getDevice('right wheel motor')

        # Adjust the low-level controller gains.
        print("%s: setting PID gains." % (self.robot_name))
        self.l_motor.setControlPID(1.0, 0.0, 0.1)
        self.r_motor.setControlPID(1.0, 0.0, 0.1)

        # Set velocity control mode.
        self.l_motor.setPosition(math.inf)
        self.r_motor.setPosition(math.inf)

        # Initialize generic behavior state machine variables.
        self.state_timer = 0        # timers in milliseconds
        self.state_index = 0        # current state

        # Thread-safe state variables.
        self.lock = threading.Lock()
        self.l_vel_target = 0.0
        self.r_vel_target = 0.0
        return

    #================================================================
    # Thread-safe motion primitives.
    def go_forward(self, velocity):
        """Command the motor to turn at the rate which produce the ground velocity
           specified in meters/sec.  Negative values turn backward. """
        print("Starting forward motion at %f m/s." % (velocity))

        # velocity control mode
        self.l_motor.setPosition(math.inf)
        self.r_motor.setPosition(math.inf)

        # calculate the rotational rate in radians/sec based on the wheel radius
        theta_dot = velocity / self.wheel_radius
        with self.lock:
            self.l_vel_target = theta_dot
            self.r_vel_target = theta_dot
        return

    def go_rotate(self, rot_velocity):
        """Command the motors to turn in place at the rate which produce the rotational
           velocity specified in radians/sec.  Negative values turn
           backward."""

        print("Starting rotation at %f rad/s." % (rot_velocity))

        # calculate the difference in linear velocity of the wheels
        linear_velocity = self.axle_length * rot_velocity

        # calculate the net rotational rate in radians/sec based on the wheel radius
        theta_dot = linear_velocity / self.wheel_radius

        # apply the result symmetrically to the wheels
        with self.lock:
            self.l_vel_target = 0.5*theta_dot
            self.r_vel_target = -0.5*theta_dot

        return

    #================================================================
    def run(self):
        # Run loop to execute a periodic script until the simulation quits.
        # If the controller returns -1, the simulator is quitting.
        while self.step(EVENT_LOOP_DT) != -1:
            # Read simulator clock time.
            self.sim_time = self.getTime()

            # Update the actuator controller targets.
            with self.lock:
                self.l_motor.setVelocity(self.l_vel_target)
                self.r_motor.setVelocity(self.r_vel_target)


################################################################
# Start the script.
robot = Wobbly()
script = Script(robot)
script.start()
robot.run()