from controller import Robot, Motor, DistanceSensor, Compass, GPS
import math

# Constants
MAX_SPEED = 6.28

# Create the Robot instance.
robot = Robot()

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

# Get devices
left_motor = robot.getDevice("left wheel motor")
right_motor = robot.getDevice("right wheel motor")
ball_sensor = robot.getDevice("ball sensor") # IR ball sensor
compass = robot.getDevice("compass")
gps = robot.getDevice("gps")

# Try to get sonar sensors for obstacle avoidance
sonar_names = ["sonar0", "sonar1", "sonar2", "sonar3", "sonar4", "sonar5", "sonar6", "sonar7"]
sonars = []
for i, name in enumerate(sonar_names):
    try:
        sensor = robot.getDevice(name)
        sensor.enable(timestep)
        sonars.append(sensor)
    except: 
        print(f"Warning: Sonar device \'{name}\' not found.")
        pass

# Enable devices
ball_sensor.enable(timestep)
compass.enable(timestep)
gps.enable(timestep)

# Set motor positions to infinity (velocity control)
left_motor.setPosition(float("inf"))
right_motor.setPosition(float("inf"))

# Function to set motor speeds
def set_speeds(left_speed, right_speed):
    left_motor.setVelocity(left_speed)
    right_motor.setVelocity(right_speed)

# Define goal positions (example values, adjust based on your Webots world)
# Assuming blue team is on the left, yellow on the right
# For blue robot, opponent goal is on the right (positive X)
# For yellow robot, opponent goal is on the left (negative X)
# Let's assume our robot is on the blue team for now (facing positive X)
OPPONENT_GOAL_X = 0.7 # Example X coordinate of opponent's goal
OWN_GOAL_X = -0.7 # Example X coordinate of own goal

# Main loop:
# Perform simulation steps until Webots is stopped.
while robot.step(timestep) != -1:
    # Read ball sensor value
    ball_direction = ball_sensor.getValue()
    
    # Get robot's current position and orientation
    robot_position = gps.getValues()
    # compass_values = compass.getValues()
    # For more advanced orientation, you would use compass_values and math.atan2
    # to get the robot's current heading in radians.

    # Obstacle avoidance logic (if sonars are available)
    obstacle_detected = False
    for sonar in sonars:
        if sonar.getValue() < 0.5: # If an obstacle is detected within 0.5 meters
            obstacle_detected = True
            break

    if obstacle_detected:
        # Simple avoidance: turn away from the obstacle
        set_speeds(-MAX_SPEED, MAX_SPEED) # Turn left to avoid
    elif ball_direction[2] > 0.05: # Ball is detected and in front (robot has 'possession' or is close to ball)
        # Robot has the ball, now decide to attack or defend based on position
        
        # Determine if robot is in a defensive or offensive zone
        # This is a simplified zone definition, adjust as needed
        if robot_position[0] < -0.3: # Defensive zone (closer to own goal)
            # Prioritize clearing the ball or moving it towards center
            # If ball is directly in front, push it forward
            turn_speed = ball_direction[0] * MAX_SPEED * 0.3 # Less aggressive turning for defense
            forward_speed = MAX_SPEED - abs(turn_speed)
            set_speeds(forward_speed - turn_speed, forward_speed + turn_speed)
        elif robot_position[0] > 0.3: # Offensive zone (closer to opponent's goal)
            # Prioritize scoring
            turn_speed = ball_direction[0] * MAX_SPEED * 0.7 # More aggressive turning for offense
            forward_speed = MAX_SPEED - abs(turn_speed)
            set_speeds(forward_speed - turn_speed, forward_speed + turn_speed)
        else: # Midfield zone
            # Balance between attacking and defending, primarily push towards opponent goal
            turn_speed = ball_direction[0] * MAX_SPEED * 0.5 
            forward_speed = MAX_SPEED - abs(turn_speed)
            set_speeds(forward_speed - turn_speed, forward_speed + turn_speed)

    else: # No ball detected or ball is too far/behind, implement search pattern or defensive positioning
        # If no ball, implement a more intelligent search or return to a strategic position
        
        # Option 1: Simple search pattern (rotate and move)
        # set_speeds(0.2 * MAX_SPEED, -0.2 * MAX_SPEED) # Rotate in place to search
        
        # Option 2: Move to a central defensive position if ball is lost
        # This requires knowing the field dimensions and a target central point
        target_x = 0.0 # Center of the field X coordinate
        target_z = 0.0 # Center of the field Z coordinate
        
        # Calculate distance to target
        distance_to_target = math.sqrt((robot_position[0] - target_x)**2 + (robot_position[2] - target_z)**2)
        
        if distance_to_target > 0.1: # If not at target, move towards it
            # Calculate angle to target (simplified, assumes robot is facing positive X)
            # This would need compass for accurate heading
            # For now, just move towards center if far
            if robot_position[0] < target_x: # If to the left of center
                set_speeds(MAX_SPEED * 0.5, MAX_SPEED * 0.5) # Move forward (towards positive X)
            else: # If to the right of center
                set_speeds(-MAX_SPEED * 0.5, -MAX_SPEED * 0.5) # Move backward (towards negative X)
        else: # At target, rotate to search
            set_speeds(0.2 * MAX_SPEED, -0.2 * MAX_SPEED) # Rotate in place to search

    pass