# Robot 1 – Soccer controller (procedural)
# Author: Hosam Yasin
# Notes: uses only built-in Webots devices (no external libs).

from controller import Robot
from math import atan2, cos, sin, pi

# ---------------- basic setup ----------------
robot = Robot()
TIME_STEP = int(robot.getBasicTimeStep())

# try to find a device by trying several common names
def pick(robot, names):
    for n in names:
        try:
            d = robot.getDevice(n)
            if d:
                return d
        except:
            pass
    return None

# motors
left_motor  = pick(robot, ["left wheel motor", "left_motor", "motor_left"])
right_motor = pick(robot, ["right wheel motor", "right_motor", "motor_right"])
for m in (left_motor, right_motor):
    if m:
        m.setPosition(float('inf'))
        m.setVelocity(0.0)

# sensors
compass = pick(robot, ["compass", "magnetometer"])
gps     = pick(robot, ["gps"])

if compass: compass.enable(TIME_STEP)
if gps:     gps.enable(TIME_STEP)

# distance sensors (for avoidance)
dist = []
for prefix in ["ps", "ds", "so", "us"]:
    for i in range(16):
        d = pick(robot, [f"{prefix}{i}"])
        if d:
            d.enable(TIME_STEP)
            dist.append((f"{prefix}{i}", d))

# IR ring for ball (names vary a lot; we’ll try common patterns)
ir = []
for prefix in ["ir", "irSeeker", "irs", "ball_ir"]:
    for i in range(16):
        s = pick(robot, [f"{prefix}{i}"])
        if s:
            s.enable(TIME_STEP)
            ir.append((f"{prefix}{i}", s))

# --------------- helpers ---------------
MAX_SPEED    = 7.0
BASE_SPEED   = 4.8
CHASE_SPEED  = 6.6
TURN_GAIN    = 3.0
ALIGN_GAIN   = 2.2
ALIGN_TOL    = 15 * pi/180
SIDE_NEAR    = 0.25
FRONT_NEAR   = 0.30
BALL_MIN     = 0.02
BALL_CLOSE   = 0.25
SEARCH_DIF   = 1.3

def clamp(x, lo, hi):
    return lo if x < lo else (hi if x > hi else x)

def drive(l, r):
    if left_motor:  left_motor.setVelocity(clamp(l, -MAX_SPEED, MAX_SPEED))
    if right_motor: right_motor.setVelocity(clamp(r, -MAX_SPEED, MAX_SPEED))

def read_prox():
    # return (front, left, right) as simple maxima
    front = left = right = 0.0
    for name, s in dist:
        try:
            v = float(s.getValue())
        except:
            v = 0.0
        # rough grouping by index if available
        idx = -1
        num = ''.join([c for c in name if c.isdigit()])
        if num != "": idx = int(num)
        if idx in [2,3,4,5] or idx == -1:
            front = max(front, v)
        if idx in [6,7]:
            left = max(left, v)
        if idx in [0,1]:
            right = max(right, v)
    return front, left, right

def ball_bearing():
    if not ir:
        return None, 0.0
    sx = sy = total = 0.0
    n = len(ir)
    for i, (_, s) in enumerate(ir):
        try:
            val = s.getValue()
        except:
            val = 0.0
        # simple normalization
        w = val/1000.0 if val > 1.0 else val
        ang = 2*pi*i/n
        sx += w*cos(ang)
        sy += w*sin(ang)
        total += w
    if total < BALL_MIN:
        return None, 0.0
    return atan2(sy, sx), total/n

def yaw_heading():
    if not compass:
        return None
    v = compass.getValues()  # approx
    return atan2(v[0], v[2])

def goal_direction():
    if not gps:
        return 1.0, 0.0
    x, _, z = gps.getValues()
    # determine side from initial reading
    return (-1.0 if init_x > 0 else 1.0) - 0.0*x, -z

def norm(a):
    while a > pi:  a -= 2*pi
    while a < -pi: a += 2*pi
    return a

# warm up gps a few steps to get initial side
init_x = 0.0
if gps:
    for _ in range(3): robot.step(TIME_STEP)
    init_x = gps.getValues()[0]

# --------------- simple FSM ---------------
SEARCH, CHASE, ALIGN, AVOID = 0, 1, 2, 3
state = SEARCH
burst = 0
BURST_STEPS = 20

while robot.step(TIME_STEP) != -1:
    ba, bs = ball_bearing()
    f, l, r = read_prox()
    front_block = f >= FRONT_NEAR
    side_block  = (l >= SIDE_NEAR) or (r >= SIDE_NEAR)

    if state == SEARCH:
        if ba is not None:
            state = CHASE
        drive(+SEARCH_DIF, -SEARCH_DIF)

    elif state == CHASE:
        if front_block or side_block:
            state = AVOID
        if ba is None:
            state = SEARCH
            continue
        turn = TURN_GAIN * norm(ba)
        fwd  = BASE_SPEED
        if bs >= BALL_CLOSE:
            state = ALIGN
            burst = 0
        drive(fwd - turn, fwd + turn)

    elif state == ALIGN:
        if front_block:
            state = AVOID
            burst = 0
        dx, dz = goal_direction()
        head = yaw_heading()
        if head is None:
            err = norm(ba if ba is not None else 0.0)
        else:
            desired = atan2(dx, -dz)
            err = norm(desired - head)
        if abs(err) < ALIGN_TOL:
            if burst < BURST_STEPS:
                drive(CHASE_SPEED, CHASE_SPEED)
                burst += 1
            else:
                state = SEARCH
        else:
            turn = ALIGN_GAIN * err
            drive(BASE_SPEED - turn, BASE_SPEED + turn)

    elif state == AVOID:
        if front_block:
            delta = (l - r)
            drive(-BASE_SPEED - delta, -BASE_SPEED + delta)
        else:
            delta = (l - r)
            drive(BASE_SPEED - delta, BASE_SPEED + delta)
        if not front_block and not side_block:
            state = CHASE if ba is not None else SEARCH

    else:
        state = SEARCH