# Advanced Intelligent Robot 1
from controller import Robot
import math

TIME_STEP = 64
MAX_SPEED = 6.28  # Maximum wheel speed

# Initialize robot
robot = Robot()

# Wheels
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)

# Sensors
ir_sensor = robot.getDevice('ir_ball_sensor')
ir_sensor.enable(TIME_STEP)

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

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

sonar = robot.getDevice('sonar')
sonar.enable(TIME_STEP)

# Goal position (adjust based on team: Yellow or Blue)
GOAL_X = 0.0  # Example goal X coordinate
GOAL_Y = 1.0  # Example goal Y coordinate

# Helper functions
def get_robot_position():
    pos = gps.getValues()
    return pos[0], pos[1]

def get_ball_direction():
    # Use IR sensor to detect ball
    value = ir_sensor.getValue()
    if value > 0.5:
        return 0  # Ball is in front
    else:
        return 1  # Need to rotate to search

def get_heading_to_goal():
    rx, ry = get_robot_position()
    dx = GOAL_X - rx
    dy = GOAL_Y - ry
    angle = math.atan2(dy, dx)
    return angle

def rotate_towards(angle):
    # Rotate robot towards a specific angle
    north = compass.getValues()
    heading = math.atan2(north[1], north[0])
    error = angle - heading
    # Adjust wheel speeds based on error
    left_speed = MAX_SPEED * (1 - 0.5 * error)
    right_speed = MAX_SPEED * (1 + 0.5 * error)
    left_motor.setVelocity(max(min(left_speed, MAX_SPEED), -MAX_SPEED))
    right_motor.setVelocity(max(min(right_speed, MAX_SPEED), -MAX_SPEED))

def move_forward():
    left_motor.setVelocity(MAX_SPEED)
    right_motor.setVelocity(MAX_SPEED)

def avoid_obstacles():
    if sonar.getValue() < 0.5:
        # Quick turn to avoid obstacles
        left_motor.setVelocity(-MAX_SPEED / 2)
        right_motor.setVelocity(MAX_SPEED / 2)
        return True
    return False

def play():
    # If obstacle detected, avoid first
    if avoid_obstacles():
        return
    
    # Follow the ball
    if ir_sensor.getValue() > 0.5:
        # Ball in front
        goal_angle = get_heading_to_goal()
        rotate_towards(goal_angle)
        move_forward()
    else:
        # Ball not detected, rotate to search
        left_motor.setVelocity(MAX_SPEED / 2)
        right_motor.setVelocity(-MAX_SPEED / 2)

# Main loop
while robot.step(TIME_STEP) != -1:
    play()