# boom_monopod
# Copyright, 2024, Garth Zeglin.
# Sample Webots controller file for driving a monopod mounted on a passive boom.

print("boom_monopod.py waking up.")

# Import standard Python libraries.
import math

# Import the Webots simulator API.
from controller import Robot

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

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

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

        # Connect to the radio receiver.
        self.receiver = self.getDevice('receiver')
        self.receiver.enable(100) # milliseconds of sampling period
        self.message_base_address = "/" + self.robot_name

        # Fetch handles for the motors.
        self.pitch_motor = self.getDevice('pitch')
        self.leg_actuator = self.getDevice('leg_actuator')

        # Use the leg linear actuator in force mode.
        self.leg_actuator.setPosition(math.inf)
        self.leg_actuator.setForce(0)

        # The controller uses a fixed thrust on takeoff with no altitude control.
        # The units are meters of leg spring compression.
        self.thrust = 0.040

        # Target speed in boom_z radians/sec.
        self.speed = 1.0

        # Fetch handles for the joint sensors.
        self.boom_z  = self.getDevice('boom_z')
        self.boom_y  = self.getDevice('boom_y')
        self.boom_x  = self.getDevice('boom_x')
        self.leg_len = self.getDevice('leg_encoder')

        # Enable the sensors at full cycle rate.
        for s in [self.boom_z, self.boom_y, self.boom_x, self.leg_len]:
            s.enable(EVENT_LOOP_DT)

        # Initialize state estimation.
        self.delta_t = EVENT_LOOP_DT * 0.001
        self.sim_time = 0.0
        self.b_x, self.b_y, self.b_z    = (0.0, 0.0, 0.0)
        self.b_xd, self.b_yd, self.b_zd = (0.0, 0.0, 0.0)
        self.b_x_last, self.b_y_last, self.b_z_last = (0.0, 0.0, 0.0)
        self.leg_length = 0.0

        # The pitch motor is driven in position mode.
        self.pitch_motor.setPosition(0.0)

    #================================================================
    def poll_sensors(self):
        """Read sensors and update filters."""

        # Read the boom angles to determine the spatial position.  A positive b_z is
        # forward travel, a positive b_y is foot to ground contact, a positive b_x
        # is negative pitch, i.e. foot placed forward.
        self.b_z = self.boom_z.getValue()
        self.b_y = self.boom_y.getValue()
        self.b_x = self.boom_x.getValue()

        # Calculate velocities using finite difference.
        self.b_zd = (self.b_z - self.b_z_last) / self.delta_t
        self.b_yd = (self.b_y - self.b_y_last) / self.delta_t
        self.b_xd = (self.b_x - self.b_x_last) / self.delta_t

        self.b_x_last, self.b_y_last, self.b_z_last = (self.b_x, self.b_y, self.b_z)

        # Read the leg length sensor as a step toward simulating a leg spring.  Leg
        # length is positive as the leg extends, i.e., is normally decreasing during
        # landing and increasing during takeoff.
        self.leg_length = self.leg_len.getValue()
        return

    #================================================================
    def poll_control(self):
        """Update control model.  Note that the leg spring is simulated using a
        force actuator, so this also computes the passive leg spring physics."""

        # Choose a pitch angle proportional to the speed error.
        self.pitch_motor.setPosition(-0.3 * (self.speed - self.b_zd))

        # Superimpose thrust force during liftoff phase.  The condition
        # is that the foot is in contact but the boom velocity is upwards.
        if self.b_y > 0.0 and self.b_yd < 0.0:
            neutral = self.thrust
        else:
            neutral = 0.0

        # Apply simulated spring force.  Positive forces extend the leg, i.e., press
        # the foot against the ground.  A typical equilibrium force is about 6N
        # representing the weight of the boom.  The spring constant is chosen k = F / d
        # for roughly 50 mm of compression at equilibrium.
        spring_force = -120.0 * (self.leg_length - neutral)
        self.leg_actuator.setForce(spring_force)

        # print(f"boom: y:{b_y}, z:{b_z}  leg_length: {leg_length} neutral: {neutral} force: {spring_force}")
        return

   #================================================================
    def poll_receiver(self):
        """Process all available radio messages."""
        while self.receiver.getQueueLength() > 0:
            packet = self.receiver.getString()
            tokens = packet.split()
            if len(tokens) >= 2:
                if tokens[0].startswith(self.message_base_address):
                    if tokens[0] == self.message_base_address + '/speed':
                        self.speed = float(tokens[1])
                        print("%s: set speed to %f" % (robot_name, self.speed))
            # done with packet processing, advance to the next packet
            self.receiver.nextPacket()

   #================================================================
    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()

            # Process sensor values.
            self.poll_sensors()

            # Poll the simulated radio receiver.
            self.poll_receiver()

            # Update control.
            self.poll_control()


################################################################
# Start the script.
robot = Monopod()
robot.run()