"""
Sheep flocking controller (Webots, Reynolds boids variant).

Each sheep broadcasts its GPS position every 3 steps on channel 1 and
listens for the dog and peer sheep positions.  Peers are keyed by robot
name so each neighbour has exactly one current entry in the dict.

Force stack each step (summed then converted to a heading + speed):
    flee       — away from dog, quadratic ramp, dominant when close
    cohesion   — toward flock centre, halved while fleeing
    separation — inverse-distance push, prevents physical overlap
    walls      — linear repulsion from field boundary
    wander     — small persistent drift for natural idle motion

Pen behaviour: on first entry into the quarantine pen the sheep latches
permanently — it turns pink (via the exposed woolColor PROTO field) and
the normal force stack is replaced by pen-confinement forces only.
"""

import random
import math
from controller import Supervisor

# ---------------------------------------------------------------------------
# Tuning constants
# ---------------------------------------------------------------------------

MAX_SPEED    = 22.0   # rad/s hard clamp on both motors
FLEE_SPEED   = 20.0   # rad/s upper bound while panicking
WANDER_SPEED =  3.0   # rad/s lower bound during calm wandering

X_MIN, X_MAX = -14.5, 14.5   # stone wall inner edges (metres)
Y_MIN, Y_MAX = -14.5, 14.5
WALL_MARGIN  =  3.5           # avoidance starts this far from the wall

FLEE_DIST       = 7.0   # dog within this radius triggers flee (metres)
SEPARATION_DIST = 2.5   # inverse-distance push active inside this radius
COHESION_DIST   = 8.0   # pull toward flock centre active inside this radius

PEN_X_MIN, PEN_X_MAX = 10.0, 13.0   # quarantine pen extents (metres)
PEN_Y_MIN, PEN_Y_MAX = -15.0, -8.0  # open entrance at y=-8, gate at y=-15
PEN_MARGIN = 0.8                     # confinement force starts this far from pen wall

# ---------------------------------------------------------------------------
# Device setup
# ---------------------------------------------------------------------------

robot    = Supervisor()
timestep = int(robot.getBasicTimeStep())
name     = robot.getName()
self_node = robot.getSelf()

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

gps      = robot.getDevice("gps");      gps.enable(timestep)
compass  = robot.getDevice("compass");  compass.enable(timestep)
receiver = robot.getDevice("receiver"); receiver.enable(timestep)
emitter  = robot.getDevice("emitter")

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def norm_angle(a):
    return math.atan2(math.sin(a), math.cos(a))


def bearing():
    # Compass returns north direction in sensor frame; for this Z-up world
    # with north = +Y, atan2(n[0], n[1]) gives the standard math angle
    # (0 = east, π/2 = north) matching atan2(fy, fx) used for heading.
    n = compass.getValues()
    return math.atan2(n[0], n[1])


def drive(heading, speed):
    err = norm_angle(heading - bearing())
    # Scale forward component by cos(err): at 90° error fwd→0 so the robot
    # spins in place to realign rather than driving sideways at full speed.
    fwd = speed * max(0.0, math.cos(err))
    k = 4.0
    left_motor.setVelocity( max(-MAX_SPEED, min(MAX_SPEED, fwd - k * err)))
    right_motor.setVelocity(max(-MAX_SPEED, min(MAX_SPEED, fwd + k * err)))


def paint_pink():
    # woolColor is declared as a PROTO field with IS binding to the DEF WOOL
    # PBRAppearance baseColor.  Changing it here propagates to every USE WOOL
    # shape on the body.  Direct field access avoids PROTO-internal opacity.
    self_node.getField("woolColor").setSFColor([1.0, 0.55, 0.72])

# ---------------------------------------------------------------------------
# State
# ---------------------------------------------------------------------------

wander_angle = random.uniform(-math.pi, math.pi)
step   = 0
dog_x  = None
dog_y  = None
peers  = {}   # name → (x, y), one entry per neighbour, cleared every 30 steps
penned = False

# ---------------------------------------------------------------------------
# Main loop
# ---------------------------------------------------------------------------

while robot.step(timestep) != -1:
    step += 1
    pos = gps.getValues()
    x, y = pos[0], pos[1]

    # Pen entry: one-way latch, never unset
    if not penned and PEN_X_MIN < x < PEN_X_MAX and PEN_Y_MIN < y < PEN_Y_MAX:
        penned = True
        paint_pink()

    # Refresh peer table (clear before receiving so fresh data is never lost)
    if step % 30 == 0:
        peers.clear()
    while receiver.getQueueLength() > 0:
        msg = receiver.getString()
        receiver.nextPacket()
        p = msg.split(":")
        if p[0] == "dog" and len(p) >= 3:
            dog_x, dog_y = float(p[1]), float(p[2])
        elif p[0] == "sheep" and len(p) >= 4 and p[1] != name:
            peers[p[1]] = (float(p[2]), float(p[3]))

    fx, fy = 0.0, 0.0

    if penned:
        # Inside pen: wander freely, strong boundary forces prevent exit,
        # separation still active to avoid collisions with other penned sheep.

        pm = PEN_MARGIN
        if x < PEN_X_MIN + pm: fx += ((PEN_X_MIN + pm - x) / pm) * 15.0
        if x > PEN_X_MAX - pm: fx -= ((x - (PEN_X_MAX - pm)) / pm) * 15.0
        if y < PEN_Y_MIN + pm: fy += ((PEN_Y_MIN + pm - y) / pm) * 15.0
        if y > PEN_Y_MAX - pm: fy -= ((y - (PEN_Y_MAX - pm)) / pm) * 15.0

        for px, py in peers.values():
            dx, dy = px - x, py - y
            d = math.hypot(dx, dy)
            if 0.05 < d < SEPARATION_DIST:
                push = (SEPARATION_DIST - d) / d
                fx -= (dx / d) * push * 2.5
                fy -= (dy / d) * push * 2.5

        if random.random() < 0.02:
            wander_angle += random.uniform(-0.6, 0.6)
        fx += math.cos(wander_angle) * 0.5
        fy += math.sin(wander_angle) * 0.5

    else:
        fleeing = False

        # Flee — quadratic ramp so force grows rapidly as the dog closes in
        if dog_x is not None:
            dx   = dog_x - x
            dy   = dog_y - y
            dist = math.hypot(dx, dy)
            if 0.01 < dist < FLEE_DIST:
                fleeing = True
                t = 1.0 - dist / FLEE_DIST
                s = t * t * 20.0
                fx -= (dx / dist) * s
                fy -= (dy / dist) * s

        # Cohesion — halved while fleeing to reduce mid-panic collisions
        cx, cy, cn = 0.0, 0.0, 0
        for px, py in peers.values():
            d = math.hypot(px - x, py - y)
            if 0.3 < d < COHESION_DIST:
                cx += px; cy += py; cn += 1
        if cn > 0:
            w = 0.08 if fleeing else 0.15
            fx += (cx / cn - x) * w
            fy += (cy / cn - y) * w

        # Separation — inverse-distance: huge when nearly overlapping, fades quickly
        for px, py in peers.values():
            dx, dy = px - x, py - y
            d = math.hypot(dx, dy)
            if 0.05 < d < SEPARATION_DIST:
                push = (SEPARATION_DIST - d) / d
                fx -= (dx / d) * push * 2.5
                fy -= (dy / d) * push * 2.5

        # Walls
        if x < X_MIN + WALL_MARGIN: fx += ((X_MIN + WALL_MARGIN - x) / WALL_MARGIN) * 6.0
        if x > X_MAX - WALL_MARGIN: fx -= ((x - (X_MAX - WALL_MARGIN)) / WALL_MARGIN) * 6.0
        if y < Y_MIN + WALL_MARGIN: fy += ((Y_MIN + WALL_MARGIN - y) / WALL_MARGIN) * 6.0
        if y > Y_MAX - WALL_MARGIN: fy -= ((y - (Y_MAX - WALL_MARGIN)) / WALL_MARGIN) * 6.0

        # Wander — suppressed while fleeing so drift cannot deflect the flee heading
        if not fleeing:
            if random.random() < 0.02:
                wander_angle += random.uniform(-0.6, 0.6)
            fx += math.cos(wander_angle) * 0.5
            fy += math.sin(wander_angle) * 0.5

    # Hard-stop clamp: within 0.5 m of a wall, zero any force component that
    # would push further into it.  Prevents the flee force from pinning a sheep
    # against the boundary when the dog approaches from outside.
    HS = 0.5
    if x < X_MIN + HS and fx < 0: fx = 0.0
    if x > X_MAX - HS and fx > 0: fx = 0.0
    if y < Y_MIN + HS and fy < 0: fy = 0.0
    if y > Y_MAX - HS and fy > 0: fy = 0.0

    heading = math.atan2(fy, fx)
    mag     = math.hypot(fx, fy)
    speed   = max(WANDER_SPEED, min(FLEE_SPEED, mag * 3.0))
    drive(heading, speed)

    if step % 3 == 0:
        emitter.send(f"sheep:{name}:{x:.4f}:{y:.4f}")
