"""Sheep flocking controller (Webots).

Each sheep broadcasts its GPS position every 3 steps on channel 1 and
listens for the dog and peer sheep positions. The behavioural step is
delegated to ``herding.flocking_sim.compute_heading_speed`` so the
training environment and Webots run identical sheep dynamics.

Pen behaviour: a sheep latches to ``penned`` the first time it crosses
the south-wall gate plane into the gate corridor. Once latched it turns
pink (via the exposed ``woolColor`` PROTO field) and the force stack
switches to in-pen containment.
"""

import math
import os
import random
import sys

# --- Make the shared herding/ package importable from this controller dir ---
_HERE = os.path.dirname(os.path.abspath(__file__))
_PROJECT_ROOT = os.path.normpath(os.path.join(_HERE, "..", ".."))
if _PROJECT_ROOT not in sys.path:
    sys.path.insert(0, _PROJECT_ROOT)

from controller import Supervisor

from herding.world.diffdrive import heading_speed_to_wheels
from herding.world.flocking_sim import MAX_SPEED, compute_heading_speed
from herding.world.geometry import (
    SHEEP_MAX_WHEEL_OMEGA,
    is_penned_position,
)


# ---------------------------------------------------------------------------
# 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)
MOTOR_MAX = min(left_motor.getMaxVelocity(), SHEEP_MAX_WHEEL_OMEGA)

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 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 headings.
    n = compass.getValues()
    return math.atan2(n[0], n[1])


def drive(heading, speed_motor):
    left_w, right_w = heading_speed_to_wheels(
        heading, min(speed_motor, MAX_SPEED), bearing(), MOTOR_MAX, k_turn=4.0
    )
    left_motor.setVelocity(left_w)
    right_motor.setVelocity(right_w)


def paint_pink():
    # woolColor is declared as a PROTO field with IS binding to the DEF WOOL
    # PBRAppearance baseColor; setting it propagates to every USE WOOL shape.
    self_node.getField("woolColor").setSFColor([1.0, 0.55, 0.72])


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

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

# Stuck detection: differential-drive sheep can pin against a wall and need
# a forced reverse-and-rotate to escape. If displacement < STUCK_DIST for
# STUCK_STEPS consecutive steps, drive toward field centre.
_prev_x, _prev_y = None, None
_stuck_count = 0
STUCK_STEPS = 20
STUCK_DIST = 0.05


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

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

    # Pen entry: one-way latch. Penned sheep get pink wool and switch behaviour.
    if not penned and is_penned_position(x, y):
        penned = True
        paint_pink()

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

    dog_xy = (dog_x, dog_y) if dog_x is not None and dog_y is not None else None
    heading, speed, wander_angle = compute_heading_speed(
        x=x, y=y, penned=penned, dog_xy=dog_xy, peers=peers,
        wander_angle=wander_angle,
    )

    # Stuck detection — safety net for differential-drive wall pinning.
    if _prev_x is not None:
        moved = math.hypot(x - _prev_x, y - _prev_y)
        _stuck_count = _stuck_count + 1 if moved < STUCK_DIST else 0
    if _stuck_count >= STUCK_STEPS:
        heading = math.atan2(-y, -x)   # always points away from the boundary
        speed = MAX_SPEED
        _stuck_count = 0
    _prev_x, _prev_y = x, y

    drive(heading, speed)

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