from controller import Robot
import math

robot = Robot()
time_step = 32

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

ir = []
for i in range(5):
    sensor = robot.getDevice('ir' + str(i))
    sensor.enable(time_step)
    ir.append(sensor)


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

# GPS
gps = robot.getDevice('gps')
gps.enable(time_step)

sonar = robot.getDevice('us0')
sonar.enable(time_step)

MAX_SPEED = 6.28
FORWARD_SPEED = 5.5
TURN_SPEED = 3.0
AVOID_SPEED = 4.0
AVOID_DISTANCE = 0.45

def get_bearing_in_degrees():
    compass_values = compass.getValues()
    rad = -math.atan2(compass_values[0], compass_values[2])
    bearing = (rad + 2 * math.pi) % (2 * math.pi)
    return math.degrees(bearing)

def get_ball_direction():
    values = [sensor.getValue() for sensor in ir]
    max_val = max(values)
    if max_val < 0.15:
        return -1 
    return values.index(max_val)

def get_angle_to_goal(position, team):
    if team == "blue":
        goal_x = -0.75
    else:  # yellow
        goal_x = 0.75
    goal_y = 0
    dx = goal_x - position[0]
    dy = goal_y - position[1]
    angle = math.degrees(math.atan2(dx, dy))
    return angle

def get_team(position):
    if position[0] > 0:
        return "yellow"
    else:
        return "blue"

while robot.step(time_step) != -1:
    ball_dir = get_ball_direction()
    bearing = get_bearing_in_degrees()
    position = gps.getValues()
    distance = sonar.getValue()
    team = get_team(position)

    if distance < AVOID_DISTANCE:
        left_motor.setVelocity(-AVOID_SPEED)
        right_motor.setVelocity(AVOID_SPEED)
        continue

    if ball_dir == -1:
        left_motor.setVelocity(TURN_SPEED)
        right_motor.setVelocity(-TURN_SPEED)
        continue

    if ball_dir == 0:
        left_motor.setVelocity(TURN_SPEED)
        right_motor.setVelocity(FORWARD_SPEED)
    elif ball_dir == 1:
        left_motor.setVelocity(FORWARD_SPEED)
        right_motor.setVelocity(0.7 * FORWARD_SPEED)
    elif ball_dir == 2:
        left_motor.setVelocity(FORWARD_SPEED)
        right_motor.setVelocity(FORWARD_SPEED)
    elif ball_dir == 3:
        left_motor.setVelocity(0.7 * FORWARD_SPEED)
        right_motor.setVelocity(FORWARD_SPEED)
    elif ball_dir == 4:
        left_motor.setVelocity(FORWARD_SPEED)
        right_motor.setVelocity(TURN_SPEED)
    else:
        left_motor.setVelocity(TURN_SPEED)
        right_motor.setVelocity(-TURN_SPEED)