"""
Shepherd Dog controller (Webots, manual keyboard control).

WASD / arrow keys drive the robot.  +/- adjust speed in 10 % increments.
GPS position is broadcast every step on channel 1 so sheep controllers
can compute flee forces.  Ears wag continuously via sinusoidal position
targets — purely cosmetic.
"""

import math
from controller import Robot, Keyboard

robot    = Robot()
timestep = int(robot.getBasicTimeStep())

left_motor  = robot.getDevice("left wheel motor")
right_motor = robot.getDevice("right wheel motor")
left_motor.setPosition(float("inf"))
right_motor.setPosition(float("inf"))
left_motor.setVelocity(0.0)
right_motor.setVelocity(0.0)

lidar = robot.getDevice("lidar")
lidar.enable(timestep)
lidar.enablePointCloud()

gps     = robot.getDevice("gps");     gps.enable(timestep)
compass = robot.getDevice("compass"); compass.enable(timestep)
emitter = robot.getDevice("emitter")
receiver = robot.getDevice("receiver"); receiver.enable(timestep)

left_ear  = robot.getDevice("left ear motor")
right_ear = robot.getDevice("right ear motor")
left_ear.setPosition(float("inf"))
right_ear.setPosition(float("inf"))
left_ear.setVelocity(0.0)
right_ear.setVelocity(0.0)

keyboard = robot.getKeyboard()
keyboard.enable(timestep)

MOTOR_MAX   = left_motor.getMaxVelocity()
speed_level = 0.5   # fraction of MOTOR_MAX; adjusted by +/-

EAR_AMPLITUDE = 0.35   # rad, peak ear deflection
EAR_RATE      = 8.0    # rad/s, how fast the ears are driven
ear_phase     = 0.0

while robot.step(timestep) != -1:
    speed = MOTOR_MAX * speed_level
    turn  = speed * 0.6   # differential turn radius

    left_vel  = 0.0
    right_vel = 0.0
    key = keyboard.getKey()
    while key > 0:
        if   key in (ord('W'), Keyboard.UP):
            left_vel  = speed
            right_vel = speed
        elif key in (ord('S'), Keyboard.DOWN):
            left_vel  = -speed
            right_vel = -speed
        elif key in (ord('A'), Keyboard.LEFT):
            left_vel  = -turn
            right_vel =  turn
        elif key in (ord('D'), Keyboard.RIGHT):
            left_vel  =  turn
            right_vel = -turn
        elif key in (ord('+'), ord('=')):
            speed_level = min(1.0, speed_level + 0.1)
            print(f"Speed: {speed_level:.0%} ({MOTOR_MAX * speed_level:.1f} rad/s)")
        elif key in (ord('-'), ord('_')):
            speed_level = max(0.1, speed_level - 0.1)
            print(f"Speed: {speed_level:.0%} ({MOTOR_MAX * speed_level:.1f} rad/s)")
        key = keyboard.getKey()

    left_motor.setVelocity(left_vel)
    right_motor.setVelocity(right_vel)

    pos = gps.getValues()
    emitter.send(f"dog:{pos[0]}:{pos[1]}")

    ear_phase += 0.12
    ear_pos = EAR_AMPLITUDE * math.sin(ear_phase)
    left_ear.setVelocity(EAR_RATE)
    right_ear.setVelocity(EAR_RATE)
    left_ear.setPosition( ear_pos)
    right_ear.setPosition(-ear_pos)
