# Modified Robot 1 – Final Robot Soccer Competition
# Author: Saleh Bin Laden
# Notes:
# Enhanced version with adaptive turning and stable kick logic for improved control.

from controller import Robot
import math

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

LEFT_MOTOR_NAME  = 'left wheel motor'
RIGHT_MOTOR_NAME = 'right wheel motor'
COMPASS_NAME     = 'compass'
GPS_NAME         = 'gps'

PROX_NAMES = [f'ps{i}' for i in range(8)]
BALL_NAMES = [f'ir{i}' for i in range(8)]

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]

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)

MAX_SPEED = 6.28
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 ----------
FIELD_X_HALF = 0.75
FIELD_Y_HALF = 0.55
SAFE_MARGIN  = 0.08
FRONT_OBS_TH = 80.0
SIDE_OBS_TH  = 60.0
BALL_SEEN_TH = 30.0
BALL_NEAR_TH = 220.0
KP_TURN    = 2.0
KV_FORWARD = 0.5
SEARCH_SPEED = 0.4 * MAX_SPEED
KICK_TIME_STEPS = int(0.35 * 1000 / TIME_STEP)

attack_dir_sign = +1
spawn_read = False
kick_timer = 0
stable_frames = 0

ir_angles = [0.0, math.pi/4, math.pi/2, 3*math.pi/4, math.pi, -3*math.pi/4, -math.pi/2, -math.pi/4]

# ---------- MAIN LOOP ----------
while robot.step(TIME_STEP) != -1:
    x, y = 0.0, 0.0
    if gps:
        pos = gps.getValues()
        x, _, y = pos[0], pos[1], pos[2]
        if not spawn_read:
            attack_dir_sign = +1 if x < 0 else -1
            spawn_read = True

    heading = 0.0
    if compass:
        north = compass.getValues()
        yaw = math.atan2(north[0], north[2])
        heading = yaw + math.pi
    while heading > math.pi: heading -= 2*math.pi
    while heading < -math.pi: heading += 2*math.pi

    ps_vals = [s.getValue() if s else 0.0 for s in prox]
    front_left, front_right, left_side, right_side = ps_vals[7], ps_vals[0], ps_vals[6], ps_vals[1]

    ir_vals = [s.getValue() if s else 0.0 for s in ball]
    max_ir = max(ir_vals) if ir_vals else 0.0
    ball_detected = max_ir > BALL_SEEN_TH

    ball_angle, ball_energy = 0.0, 0.0
    if ball_detected and len(ir_vals) == 8:
        sx, sy, total = 0.0, 0.0, 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)
            ball_energy = total / 8.0
        else:
            ball_detected = False

    avoid_wall, wall_turn = False, 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
        if abs(y) > (FIELD_Y_HALF - SAFE_MARGIN):
            avoid_wall = True
            wall_turn += -0.6 if y > 0 else 0.6

    avoid_obs, obs_turn = False, 0.0
    if front_left > FRONT_OBS_TH or front_right > FRONT_OBS_TH:
        avoid_obs = True
        obs_turn = -0.9 if front_left > front_right else 0.9
    elif left_side > SIDE_OBS_TH:
        avoid_obs, obs_turn = True, 0.6
    elif right_side > SIDE_OBS_TH:
        avoid_obs, obs_turn = True, -0.6

    forward, turn = 0.0, 0.0

    if kick_timer > 0:
        forward, turn = 1.0, 0.0
        kick_timer -= 1
    else:
        if ball_detected:
            turn = KP_TURN * ball_angle * (1 - 0.5 * (ball_energy / BALL_NEAR_TH))
            approach = KV_FORWARD * min(ball_energy, BALL_NEAR_TH)
            forward = approach * (1.0 - min(1.0, abs(turn)/2.5))

            facing_attack = True
            if compass:
                desired_heading = 0.0 if attack_dir_sign > 0 else math.pi
                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)

            if ball_energy > BALL_NEAR_TH * 0.85 and abs(ball_angle) < (18.0 * math.pi / 180.0) and facing_attack:
                stable_frames += 1
            else:
                stable_frames = 0

            if stable_frames > 3:
                kick_timer = KICK_TIME_STEPS
                stable_frames = 0
        else:
            forward, turn = 0.05, (-0.6 if attack_dir_sign < 0 else 0.6)

    if abs(x) > FIELD_X_HALF - 2*SAFE_MARGIN or abs(y) > FIELD_Y_HALF - 2*SAFE_MARGIN:
        forward *= 0.5

    if avoid_wall or avoid_obs:
        forward *= 0.3
        pri = 0.0
        if avoid_wall: pri += wall_turn
        if avoid_obs: pri += obs_turn * 0.8
        turn = 0.7 * pri + 0.3 * turn

    v_l = (forward - 0.5*turn) * MAX_SPEED
    v_r = (forward + 0.5*turn) * MAX_SPEED
    v_l = max(-MAX_SPEED, min(MAX_SPEED, v_l))
    v_r = max(-MAX_SPEED, min(MAX_SPEED, v_r))

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