TIR_PROJ/herding/world/geometry.py

"""World geometry and robot specs.

Coordinates are metres; (0, 0) is the field centre, +x east, +y north.
These constants mirror ``worlds/field.wbt`` and the proto files — if
the world changes, this file is the single point of update.

    field (rectangular)
    +-----------+
    |           |
    |  ......   |
    +---||||----+   y = -15  (south wall, 3 m gate at x in [10, 13])
        ||||
        |pen|       y in [-22, -15]
        +---+

    field_round (circular, R = 15 m)
         .---.
       /  ...  \\
      |  .....  |    gate at south, x in [-1.83, 1.83]
       \\  ...  /
         '-+-'       pen y in [-22, -15]
"""

import os
import math

# ---------------------------------------------------------------------------
# Field shape selection — controlled by HERDING_WORLD env var at runtime.
# Defaults to "field" (rectangular). The launcher writes it into the
# runtime cfg so the controller can pick it up too.
# ---------------------------------------------------------------------------
FIELD_SHAPE = (os.environ.get("HERDING_WORLD", "field")).lower()


# ==================== Rectangular field (field.wbt) ====================
FIELD_X = (-15.0, 15.0)
FIELD_Y = (-15.0, 15.0)
FIELD_INSIDE_MARGIN = 0.5

# Pen (external, south of the field)
PEN_X = (10.0, 13.0)
PEN_Y = (-22.0, -15.0)
PEN_CENTER = (0.5 * (PEN_X[0] + PEN_X[1]), 0.5 * (PEN_Y[0] + PEN_Y[1]))
PEN_ENTRY = (0.5 * (PEN_X[0] + PEN_X[1]), -15.0)

# Gate (hole in the south wall)
GATE_X = PEN_X
GATE_Y = -15.0


# ==================== Round field (field_round.wbt) ====================
FIELD_ROUND_R = 15.0
FIELD_ROUND_PEN_X = (-1.5, 1.5)
FIELD_ROUND_PEN_Y = (-22.0, -15.0)
FIELD_ROUND_PEN_CENTER = (
    0.5 * (FIELD_ROUND_PEN_X[0] + FIELD_ROUND_PEN_X[1]),
    0.5 * (FIELD_ROUND_PEN_Y[0] + FIELD_ROUND_PEN_Y[1]),
)
FIELD_ROUND_PEN_ENTRY = (0.0, -15.0)
FIELD_ROUND_GATE_X = FIELD_ROUND_PEN_X
FIELD_ROUND_GATE_Y = -15.0


# ==================== Active geometry (resolved at import) ===============
# Rectangular defaults are already assigned above.  Override for round.
if FIELD_SHAPE == "field_round":
    PEN_X = FIELD_ROUND_PEN_X
    PEN_Y = FIELD_ROUND_PEN_Y
    PEN_CENTER = FIELD_ROUND_PEN_CENTER
    PEN_ENTRY = FIELD_ROUND_PEN_ENTRY
    GATE_X = FIELD_ROUND_GATE_X
    GATE_Y = FIELD_ROUND_GATE_Y


def configure_from_args(argv: list[str] | None = None) -> str:
    """Parse ``--world`` from *argv* (or ``sys.argv[1:]``), call :func:`configure`,
    and set ``HERDING_WORLD`` in the environment.

    Returns the resolved world name (``"field"`` or ``"field_round"``).

    Call this at the very top of any script that accepts a ``--world`` flag,
    *before* importing anything from ``herding.*`` that depends on field
    geometry.  This centralises the pre-parse logic that was previously
    duplicated in ``bc/collect.py``, ``rl/train.py``, and ``eval.py``::

        from herding.world.geometry import configure_from_args
        configure_from_args()   # reads sys.argv
    """
    import os
    import sys as _sys
    args = argv if argv is not None else _sys.argv[1:]
    world = "field"
    for i, a in enumerate(args):
        if a == "--world" and i + 1 < len(args):
            world = args[i + 1]
            break
        if a.startswith("--world="):
            world = a.split("=", 1)[1]
            break
    configure(world)
    os.environ["HERDING_WORLD"] = world
    return world


def configure(shape: str) -> None:
    """Switch the active field geometry at runtime.

    Call this **before** importing any other ``herding.*`` module that
    depends on the constants below (flocking_sim, lidar_sim, obs, etc.).
    The import-time env-var path (``HERDING_WORLD``) still works; this
    function is for scripts that need to choose the world via a CLI flag.
    """
    global FIELD_SHAPE, PEN_X, PEN_Y, PEN_CENTER, PEN_ENTRY, GATE_X, GATE_Y
    shape = shape.lower()
    FIELD_SHAPE = shape
    if shape == "field_round":
        PEN_X = FIELD_ROUND_PEN_X
        PEN_Y = FIELD_ROUND_PEN_Y
        PEN_CENTER = FIELD_ROUND_PEN_CENTER
        PEN_ENTRY = FIELD_ROUND_PEN_ENTRY
        GATE_X = FIELD_ROUND_GATE_X
        GATE_Y = FIELD_ROUND_GATE_Y
    else:
        PEN_X = (10.0, 13.0)
        PEN_Y = (-22.0, -15.0)
        PEN_CENTER = (0.5 * (PEN_X[0] + PEN_X[1]), 0.5 * (PEN_Y[0] + PEN_Y[1]))
        PEN_ENTRY = (0.5 * (PEN_X[0] + PEN_X[1]), -15.0)
        GATE_X = PEN_X
        GATE_Y = -15.0

# Dog spec — protos/ShepherdDog.proto
DOG_WHEEL_RADIUS = 0.038         # m
DOG_WHEEL_BASE = 0.28            # m, axle-to-axle
DOG_MAX_WHEEL_OMEGA = 70.0       # rad/s
DOG_MAX_LINEAR = DOG_WHEEL_RADIUS * DOG_MAX_WHEEL_OMEGA  # ≈ 2.66 m/s

# Dog mecanum spec — 4-wheel omnidirectional layout
DOG_WHEEL_BASE_X = 0.28            # m, front-to-back axle distance
DOG_WHEEL_BASE_Y = 0.28            # m, left-to-right axle distance

# Sheep spec — protos/Sheep.proto
SHEEP_WHEEL_RADIUS = 0.031       # m
SHEEP_WHEEL_BASE = 0.20          # m
SHEEP_MAX_WHEEL_OMEGA = 25.0     # rad/s
SHEEP_MAX_LINEAR = SHEEP_WHEEL_RADIUS * SHEEP_MAX_WHEEL_OMEGA  # ≈ 0.78 m/s

WEBOTS_DT = 0.016                # seconds (matches WorldInfo.basicTimeStep)

# Virtual south wall — env and controller both keep the dog north of this.
DOG_SOUTH_LIMIT = -14.5

MAX_SHEEP = 10


def in_pen(x: float, y: float) -> bool:
    """True if (x, y) lies inside the external pen rectangle."""
    return PEN_X[0] < x < PEN_X[1] and PEN_Y[0] < y < PEN_Y[1]


def in_field(x: float, y: float, margin: float = 0.0) -> bool:
    if FIELD_SHAPE == "field_round":
        r = FIELD_ROUND_R - margin
        return x * x + y * y <= r * r
    return (FIELD_X[0] + margin <= x <= FIELD_X[1] - margin
            and FIELD_Y[0] + margin <= y <= FIELD_Y[1] - margin)


def in_gate_corridor(x: float, y: float, margin: float = 0.0) -> bool:
    """True if (x, y) lies in the column of the gate (between field and pen)."""
    return (GATE_X[0] - margin <= x <= GATE_X[1] + margin
            and PEN_Y[0] - margin <= y <= GATE_Y + margin)


def is_penned(x: float, y: float, latch_margin: float = 0.2) -> bool:
    """True iff (x, y) is in the gate column and south of the gate line."""
    return (GATE_X[0] - latch_margin <= x <= GATE_X[1] + latch_margin
            and y <= GATE_Y)


def distance_to_pen_entry(x: float, y: float) -> float:
    return math.hypot(x - PEN_ENTRY[0], y - PEN_ENTRY[1])


def distance_to_wall(x: float, y: float) -> float:
    """Shortest distance from (x, y) to the nearest field wall.

    For a rectangular field this is the minimum Manhattan distance to the
    four bounding walls.  For a round field it is ``R - sqrt(x²+y²)``.
    Returns a negative value if the point is outside the field.
    """
    if FIELD_SHAPE == "field_round":
        return FIELD_ROUND_R - math.hypot(x, y)
    return min(
        x - FIELD_X[0], FIELD_X[1] - x,
        y - FIELD_Y[0], FIELD_Y[1] - y,
    )


def clip_to_field(x: float, y: float, margin: float = 0.2) -> tuple[float, float]:
    """Clip (x, y) inside the field boundary with a small margin.

    For round fields the point is projected radially inward if it exceeds
    the circular boundary.
    """
    if FIELD_SHAPE == "field_round":
        r = math.hypot(x, y)
        limit = FIELD_ROUND_R - margin
        if r > limit and r > 1e-6:
            scale = limit / r
            return x * scale, y * scale
        return x, y
    return (
        max(FIELD_X[0] + margin, min(FIELD_X[1] - margin, x)),
        max(FIELD_Y[0] + margin, min(FIELD_Y[1] - margin, y)),
    )