# Enhanced Robot 1: Smart movement, ball tracking, and obstacle avoidance

from controller import Robot, DistanceSensor

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

# === Initialize motors ===
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)

# === Initialize sensors ===
# Ball sensor (Infrared)
ball_sensor = robot.getDevice("ball_sensor")  # Make sure this name matches Webots setup
ball_sensor.enable(timestep)

# Front distance sensors (for obstacle avoidance)
front_left = robot.getDevice("ds_left")
front_right = robot.getDevice("ds_right")
front_left.enable(timestep)
front_right.enable(timestep)

# === Parameters ===
base_speed = 4.0
turn_speed = 2.0
avoid_distance = 80.0  # distance threshold for obstacles

# === Main control loop ===
while robot.step(timestep) != -1:
    ball_value = ball_sensor.getValue()
    left_value = front_left.getValue()
    right_value = front_right.getValue()

    # 1️⃣ If there’s an obstacle, avoid it
    if left_value < avoid_distance or right_value < avoid_distance:
        # Turn away from the obstacle
        left_motor.setVelocity(-turn_speed)
        right_motor.setVelocity(turn_speed)
        continue

    # 2️⃣ If the ball is detected → move toward it
    if ball_value > 0.1:
        # Stronger signal = closer to ball, so move straight
        left_motor.setVelocity(base_speed)
        right_motor.setVelocity(base_speed)
    else:
        # 3️⃣ If no ball detected → rotate in place to search
        left_motor.setVelocity(turn_speed)
        right_motor.setVelocity(-turn_speed)