# Modified Robot 1 – Final Robot Soccer Competition
# Author: Saleh Bin Laden
# Notes:
# - Paste this file content into robot1.py (Yellow or Blue) as instructed.
# - No external libraries used. No classes. Minimal helper logic, all inline.
# - Assumes common Webots device names; adjust the names in the "DEVICE NAMES" block if your base project uses different ones.
# - Uses IR ball ring (ir0..ir7), proximity (ps0..ps7), compass, gps, and two wheel motors.
# - Behavior: search -> approach -> align -> kick, with obstacle & wall avoidance.

from controller import Robot

# ---------- SIM & DEVICES SETUP ----------
robot = Robot()
TIME_STEP = int(robot.getBasicTimeStep()) if robot.getBasicTimeStep() > 0 else 32

# DEVICE NAMES (change if your base code differs)
LEFT_MOTOR_NAME  = 'left wheel motor'   # try 'left wheel' if needed
RIGHT_MOTOR_NAME = 'right wheel motor'  # try 'right wheel' if needed
COMPASS_NAME     = 'compass'
GPS_NAME         = 'gps'

# 8 proximity sensors around robot (e-puck style)
PROX_NAMES = [f'ps{i}' for i in range(8)]

# 8 IR ball sensors forming a ring (common in soccer kits)
BALL_NAMES = [f'ir{i}' for i in range(8)]

# Get and enable devices safely
def safe_get(name):
    try:
        return robot.getDevice(name)
    except Exception:
        return None

left_motor  = safe_get(LEFT_MOTOR_NAME)
right_motor = safe_get(RIGHT_MOTOR_NAME)
compass     = safe_get(COMPASS_NAME)
gps         = safe_get(GPS_NAME)

prox = [safe_get(n) for n in PROX_NAMES]
ball = [safe_get(n) for n in BALL_NAMES]

# Enable sensors that exist
if compass: compass.enable(TIME_STEP)
if gps: gps.enable(TIME_STEP)

for s in prox:
    if s: s.enable(TIME_STEP)

for s in ball:
    if s: s.enable(TIME_STEP)

# Configure motors for velocity control
MAX_SPEED = 6.28  # typical for e-puck; adjust if your base robot differs
if left_motor:
    left_motor.setPosition(float('inf'))
    left_motor.setVelocity(0.0)
if right_motor:
    right_motor.setPosition(float('inf'))
    right_motor.setVelocity(0.0)

# ---------- CONSTANTS (TUNE FREELY) ----------
# Field bounds (approx for Webots soccer mini-field). Used for wall avoidance via GPS.
FIELD_X_HALF = 0.75   # half-length in X (m) – adjust to your world
FIELD_Y_HALF = 0.55   # half-width  in Y (m)

SAFE_MARGIN  = 0.08   # start steering inwards when closer than this to walls

# Proximity thresholds (raw units depend on sensor model)
FRONT_OBS_TH = 80.0   # obstacle ahead
SIDE_OBS_TH  = 60.0   # obstacle on the flanks

# IR ball thresholds (raw units depend on IR board)
BALL_SEEN_TH = 30.0   # min reading to consider ball detected
BALL_NEAR_TH = 220.0  # ball is very close (go for kick)

# Motion gains
KP_TURN    = 2.0      # proportional turn gain from angle error
KV_FORWARD = 0.5      # forward speed scale vs ball strength
SEARCH_SPEED = 0.4 * MAX_SPEED  # spin speed when searching

# Kick burst
KICK_TIME_STEPS = int(0.35 * 1000 / TIME_STEP)  # ~0.35s burst

# Attack direction guess: +X by default, but infer from spawn X at t=0
attack_dir_sign = +1  # +1 means attack toward +X, -1 toward -X
spawn_read = False

kick_timer = 0

# Precompute angles for the 8 IR ball sensors (0 front, then CCW), typical mounting
# Map indices to angles in radians (0 = forward, +CCW)
import math
ir_angles = [
    0.0,                 # ir0 – front
    math.pi/4,           # ir1 – front-left
    math.pi/2,           # ir2 – left
    3*math.pi/4,         # ir3 – back-left
    math.pi,             # ir4 – back
    -3*math.pi/4,        # ir5 – back-right
    -math.pi/2,          # ir6 – right
    -math.pi/4           # ir7 – front-right
]

# ---------- MAIN LOOP ----------
while robot.step(TIME_STEP) != -1:

    # Read GPS once to infer side & avoid walls
    x = 0.0
    y = 0.0
    if gps:
        pos = gps.getValues()
        x, _, y = pos[0], pos[1], pos[2]  # Webots: x,z,y – but many worlds use x,y
        # Infer attack direction from spawn: if we start on left half, attack +X
        if not spawn_read:
            attack_dir_sign = +1 if x < 0 else -1
            spawn_read = True

    # Read compass to get absolute heading angle (radians), 0 = +X axis
    heading = 0.0
    if compass:
        north = compass.getValues()  # returns a 3D vector towards north
        # Convert to robot's yaw. Webots compass points to north; compute robot yaw w.r.t +X.
        # yaw = atan2(north_x, north_z); robot heading = yaw + pi (approx). We want 0 at +X.
        yaw = math.atan2(north[0], north[2])
        heading = yaw + math.pi  # normalize later

    # Normalize heading to [-pi, pi]
    while heading > math.pi: heading -= 2*math.pi
    while heading < -math.pi: heading += 2*math.pi

    # Read proximity sensors
    ps_vals = [s.getValue() if s else 0.0 for s in prox]
    front_left  = ps_vals[7]  # ps7 front-left (e-puck convention)
    front_right = ps_vals[0]  # ps0 front-right
    left_side   = ps_vals[6]  # ps6 left
    right_side  = ps_vals[1]  # ps1 right

    # Read ball ring and estimate direction + strength
    ir_vals = [s.getValue() if s else 0.0 for s in ball]
    max_ir = max(ir_vals) if len(ir_vals) else 0.0

    ball_detected = max_ir > BALL_SEEN_TH

    ball_angle = 0.0   # relative angle in robot frame
    ball_energy = 0.0  # aggregate strength

    if ball_detected and len(ir_vals) == 8:
        # Weighted vector sum to estimate bearing
        sx = 0.0
        sy = 0.0
        total = 0.0
        for i in range(8):
            v = ir_vals[i]
            total += v
            sx += v * math.cos(ir_angles[i])
            sy += v * math.sin(ir_angles[i])
        if total > 1e-6:
            ball_angle = math.atan2(sy, sx)  # -pi..pi
            ball_energy = total / 8.0
        else:
            ball_detected = False

    # ---------- STATE-FREE REACTIVE CONTROL ----------
    # Wall avoidance using GPS (steer inward & back off if too close)
    avoid_wall = False
    wall_turn = 0.0
    if gps and spawn_read:
        if abs(x) > (FIELD_X_HALF - SAFE_MARGIN):
            avoid_wall = True
            wall_turn = -0.8 if x * attack_dir_sign > 0 else 0.8  # turn back toward center/attack
        if abs(y) > (FIELD_Y_HALF - SAFE_MARGIN):
            avoid_wall = True
            wall_turn += -0.6 if y > 0 else 0.6

    # Obstacle avoidance using proximity
    avoid_obs = False
    obs_turn = 0.0
    if front_left > FRONT_OBS_TH or front_right > FRONT_OBS_TH:
        avoid_obs = True
        # steer away from the stronger side
        obs_turn = -0.9 if front_left > front_right else 0.9
    else:
        # side brushes
        if left_side > SIDE_OBS_TH:
            avoid_obs = True
            obs_turn = 0.6
        elif right_side > SIDE_OBS_TH:
            avoid_obs = True
            obs_turn = -0.6

    # Decide base forward & turn
    forward = 0.0
    turn = 0.0

    if kick_timer > 0:
        # Kick burst: full speed straight
        forward = 1.0
        turn = 0.0
        kick_timer -= 1
    else:
        if ball_detected:
            # Align to ball
            turn = KP_TURN * (ball_angle)
            # Approach speed scales with ball signal, but reduce if turning sharply
            approach = KV_FORWARD * min(ball_energy, BALL_NEAR_TH)
            forward = approach * (1.0 - min(1.0, abs(turn)/2.5))

            # If ball is very close and we're roughly facing attack direction, kick
            facing_attack = True
            if compass:
                # Desired global heading is towards opponent goal: 0 rad for +X, pi for -X
                desired_heading = 0.0 if attack_dir_sign > 0 else math.pi
                # Heading error (wrap)
                err = desired_heading - heading
                while err > math.pi: err -= 2*math.pi
                while err < -math.pi: err += 2*math.pi
                facing_attack = abs(err) < (20.0 * math.pi / 180.0)  # within 20 degrees

            if ball_energy > BALL_NEAR_TH * 0.85 and abs(ball_angle) < (18.0 * math.pi / 180.0) and facing_attack:
                kick_timer = KICK_TIME_STEPS
        else:
            # SEARCH behavior: spin toward attack side preference
            forward = 0.05  # slight creep
            turn = -0.6 if attack_dir_sign < 0 else 0.6

    # Blend avoidance with target behavior (avoidance has priority)
    if avoid_wall or avoid_obs:
        # damp forward when avoiding
        forward *= 0.3
        # combine turns; wall > obstacle > target
        pri = 0.0
        if avoid_wall:
            pri += wall_turn * 1.0
        if avoid_obs:
            pri += obs_turn * 0.8
        # mix
        turn = 0.7 * pri + 0.3 * turn

    # Convert forward/turn to wheel speeds
    # forward in [0..something], turn in radians -> map to differential
    # Simple mapping: v_l = (forward - turn) * MAX_SPEED scale; v_r = (forward + turn) * MAX_SPEED scale
    # Clamp safely
    v_l = (forward - 0.5*turn) * MAX_SPEED
    v_r = (forward + 0.5*turn) * MAX_SPEED

    # clip
    if v_l > MAX_SPEED: v_l = MAX_SPEED
    if v_r > MAX_SPEED: v_r = MAX_SPEED
    if v_l < -MAX_SPEED: v_l = -MAX_SPEED
    if v_r < -MAX_SPEED: v_r = -MAX_SPEED

    if left_motor:  left_motor.setVelocity(v_l)
    if right_motor: right_motor.setVelocity(v_r)