from controller import Robot

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

# Get motors
left_motor = robot.getDevice('left wheel motor')
right_motor = robot.getDevice('right wheel motor')
kicker_motor = robot.getDevice('kicker')

# Set motors to velocity control (infinite position)
left_motor.setPosition(float('inf'))
right_motor.setPosition(float('inf'))
kicker_motor.setPosition(0.0)  # Kicker starts at rest

# Set initial velocities to 0
left_motor.setVelocity(0.0)
right_motor.setVelocity(0.0)
kicker_motor.setVelocity(0.0)

# Enable sensors
gps = robot.getDevice('gps')
gps.enable(timestep)

compass = robot.getDevice('compass')
compass.enable(timestep)

# Assuming one sonar for front obstacle detection (adjust if multiple)
sonar = robot.getDevice('sonar')
sonar.enable(timestep)

# IR ball sensor (assumes it provides distance; in Webots, it might be emitter/receiver, but simplified as value)
ir_ball = robot.getDevice('ir_ball')
ir_ball.enable(timestep)

# Constants for behavior
MAX_SPEED = 6.28  # Max wheel velocity (radians/sec, adjust based on robot)
BALL_THRESHOLD = 0.5  # Distance to consider ball close for kicking (meters, tune as needed)
OBSTACLE_THRESHOLD = 0.3  # Min distance to avoid obstacles (meters)
TURN_SPEED = 2.0  # Speed for turning towards ball

# Main loop
while robot.step(timestep) != -1:
    # Read sensor values
    gps_values = gps.getValues()  # [x, y, z] position
    compass_values = compass.getValues()  # Orientation vector
    sonar_value = sonar.getValue()  # Distance to nearest obstacle
    ir_value = ir_ball.getValue()  # Ball detection (higher value = closer ball, assuming normalized 0-1 or distance)

    # Simple ball detection: if ir_value > some threshold, ball is detected
    # For simplicity, assume ir_value is distance (lower = closer); adjust based on actual sensor
    ball_distance = ir_value if ir_value > 0 else 10.0  # Default far if no detection
    ball_detected = ball_distance < 2.0  # Tune threshold for detection range


    left_vel = 0.0
    right_vel = 0.0
    kick_vel = 0.0

    if sonar_value < OBSTACLE_THRESHOLD:
        # Avoid obstacle: back up and turn right
        left_vel = -MAX_SPEED * 0.5
        right_vel = -MAX_SPEED * 0.5
        # Add slight turn to avoid
        right_vel += TURN_SPEED * 0.3
    elif ball_detected:
        # Move towards ball
        if ball_distance < BALL_THRESHOLD:
            # Close to ball: kick!
            kick_vel = 10.0  # Strong kick velocity (adjust)
            # Slow down wheels while kicking
            left_vel = MAX_SPEED * 0.3
            right_vel = MAX_SPEED * 0.3
        else:
            # Turn to face ball direction (simplified: assume forward if detected, or use compass to check goal direction)
            # For better: compute heading to goal using GPS (e.g., towards opponent's goal)
            # Assume blue team: goal at x=-1.0, use GPS x to decide direction
            target_x = -1.0 if gps_values[0] > 0 else 1.0  # Simple goal targeting based on position
            if gps_values[0] < target_x:  # Need to move left/right? Simplified forward bias
                left_vel = MAX_SPEED
                right_vel = MAX_SPEED
            else:
                # Turn towards goal
                left_vel = -TURN_SPEED
                right_vel = TURN_SPEED
    else:
        # No ball: search by spinning or moving forward slowly
        left_vel = TURN_SPEED * 0.5
        right_vel = -TURN_SPEED * 0.5  # Gentle spin search

    # Apply velocities
    left_motor.setVelocity(left_vel)
    right_motor.setVelocity(right_vel)
    kicker_motor.setVelocity(kick_vel)

# End of controller (loop exits on simulation end)