from controller import Robot
import math

# Create the Robot instance
robot = Robot()

# Get the time step of the current world
timestep = int(robot.getBasicTimeStep())

# Enable devices
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)

# Enable sensors
ball_sensor = robot.getDevice('ball sensor')
ball_sensor.enable(timestep)

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

sonar_left = robot.getDevice('sonar left')
sonar_right = robot.getDevice('sonar right')
sonar_left.enable(timestep)
sonar_right.enable(timestep)

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

# State variables
BALL_DETECTION_THRESHOLD = 0.1
OBSTACLE_DISTANCE = 0.3
MAX_SPEED = 6.28

# Field boundaries (adjust based on your field size)
FIELD_MIN_X = -0.5
FIELD_MAX_X = 0.5
FIELD_MIN_Y = -0.3
FIELD_MAX_Y = 0.3

def normalize_angle(angle):
    """Normalize angle to be between -pi and pi"""
    while angle > math.pi:
        angle -= 2.0 * math.pi
    while angle < -math.pi:
        angle += 2.0 * math.pi
    return angle

def get_ball_direction():
    """Get ball direction from sensor and return angle"""
    ball_values = ball_sensor.getValues()
    if ball_values[0] == 0 and ball_values[1] == 0 and ball_values[2] == 0:
        return None  # No ball detected
    
    # Calculate angle to ball
    angle = math.atan2(ball_values[0], ball_values[2])
    return angle

def get_robot_heading():
    """Get robot heading from compass"""
    compass_values = compass.getValues()
    heading = math.atan2(compass_values[0], compass_values[1])
    return heading

def get_obstacle_distance():
    """Get distance to obstacles from sonars"""
    left_dist = sonar_left.getValue()
    right_dist = sonar_right.getValue()
    return left_dist, right_dist

def get_robot_position():
    """Get robot position from GPS"""
    gps_values = gps.getValues()
    return gps_values[0], gps_values[1]

def is_in_field(x, y):
    """Check if position is within field boundaries"""
    return FIELD_MIN_X <= x <= FIELD_MAX_X and FIELD_MIN_Y <= y <= FIELD_MAX_Y

def avoid_obstacles(left_dist, right_dist):
    """Avoid obstacles based on sonar readings"""
    if left_dist < OBSTACLE_DISTANCE or right_dist < OBSTACLE_DISTANCE:
        if left_dist < right_dist:
            # Obstacle on left, turn right
            return -0.5 * MAX_SPEED, 0.5 * MAX_SPEED
        else:
            # Obstacle on right, turn left
            return 0.5 * MAX_SPEED, -0.5 * MAX_SPEED
    return None, None

def move_towards_ball(ball_angle):
    """Move towards the ball"""
    if ball_angle is None:
        # No ball detected, search for it
        return 0.3 * MAX_SPEED, -0.3 * MAX_SPEED  # Rotate in place
    
    # Move towards ball with proportional control
    gain = 2.0
    left_speed = MAX_SPEED - gain * ball_angle
    right_speed = MAX_SPEED + gain * ball_angle
    
    # Limit speeds to maximum
    left_speed = max(min(left_speed, MAX_SPEED), -MAX_SPEED)
    right_speed = max(min(right_speed, MAX_SPEED), -MAX_SPEED)
    
    return left_speed, right_speed

def stay_in_field(x, y):
    """Ensure robot stays within field boundaries"""
    if not is_in_field(x, y):
        # Move towards center
        center_x = (FIELD_MIN_X + FIELD_MAX_X) / 2
        center_y = (FIELD_MIN_Y + FIELD_MAX_Y) / 2
        
        angle_to_center = math.atan2(center_y - y, center_x - x)
        current_heading = get_robot_heading()
        
        angle_diff = normalize_angle(angle_to_center - current_heading)
        
        gain = 3.0
        left_speed = MAX_SPEED - gain * angle_diff
        right_speed = MAX_SPEED + gain * angle_diff
        
        return left_speed, right_speed
    return None, None

# Main control loop
while robot.step(timestep) != -1:
    # Get sensor readings
    ball_angle = get_ball_direction()
    left_dist, right_dist = get_obstacle_distance()
    robot_x, robot_y = get_robot_position()
    
    # Default speeds
    left_speed = 0.0
    right_speed = 0.0
    
    # Priority 1: Avoid obstacles
    avoid_left, avoid_right = avoid_obstacles(left_dist, right_dist)
    if avoid_left is not None and avoid_right is not None:
        left_speed, right_speed = avoid_left, avoid_right
    
    # Priority 2: Stay in field
    elif stay_left, stay_right := stay_in_field(robot_x, robot_y):
        left_speed, right_speed = stay_left, stay_right
    
    # Priority 3: Track and move towards ball
    else:
        left_speed, right_speed = move_towards_ball(ball_angle)
    
    # Apply motor speeds
    left_motor.setVelocity(left_speed)
    right_motor.setVelocity(right_speed)