Marcus/API/zmq_api.py

"""
zmq_api.py — ZMQ velocity + command interface to Holosoma

Previously the PUB socket was bound at module import time. That made the
module unsafe to re-import from any multiprocessing child (e.g. the LiDAR
SLAM_worker spawn), because the child would try to rebind the same port
and crash with `Address already in use`.

The bind now lives in init_zmq() — call it once from the brain entrypoint.
Child processes can import this module without any network side effects.
"""
import json
import os
import time
import zmq
from Core.config_loader import load_config
from Core.logger import log

_cfg = load_config("ZMQ")

ZMQ_HOST        = _cfg["zmq_host"]
ZMQ_PORT        = _cfg["zmq_port"]
STOP_ITERATIONS = _cfg["stop_iterations"]
STOP_DELAY      = _cfg["stop_delay"]
STEP_PAUSE      = _cfg["step_pause"]

# Shared state. These stay None until init_zmq() is called.
ctx: zmq.Context = None
sock: zmq.Socket = None
_INIT_SETTLE = 0.5       # seconds to let PUB tell subscribers it's alive


def init_zmq() -> zmq.Socket:
    """
    Bind the PUB socket. Idempotent — safe to call more than once.
    Call from the main (parent) process only. Do NOT call from multiprocessing
    children — they inherit nothing useful from the bound socket anyway.
    """
    global ctx, sock
    if sock is not None:
        return sock
    ctx  = zmq.Context()
    sock = ctx.socket(zmq.PUB)
    sock.bind(f"tcp://{ZMQ_HOST}:{ZMQ_PORT}")
    time.sleep(_INIT_SETTLE)
    log(f"ZMQ PUB bound on tcp://{ZMQ_HOST}:{ZMQ_PORT} (pid={os.getpid()})",
        "info", "zmq")
    return sock


def _ensure_sock() -> zmq.Socket:
    if sock is None:
        raise RuntimeError(
            "zmq_api not initialized — call init_zmq() from the brain "
            "entrypoint before using send_vel/send_cmd/gradual_stop"
        )
    return sock


def get_socket():
    """Return the shared ZMQ PUB socket (for odometry to reuse)."""
    return _ensure_sock()


def send_vel(vx: float = 0.0, vy: float = 0.0, vyaw: float = 0.0):
    """Send velocity to Holosoma. vx m/s | vy m/s | vyaw rad/s"""
    _ensure_sock().send_string(json.dumps({"vel": {"vx": vx, "vy": vy, "vyaw": vyaw}}))


def gradual_stop():
    """Smooth deceleration to zero over ~1 second."""
    s = _ensure_sock()
    for _ in range(STOP_ITERATIONS):
        s.send_string(json.dumps({"vel": {"vx": 0.0, "vy": 0.0, "vyaw": 0.0}}))
        time.sleep(STOP_DELAY)


def send_cmd(cmd: str):
    """Send Holosoma state command: start | walk | stand | stop"""
    _ensure_sock().send_string(json.dumps({"cmd": cmd}))


# Load navigation constants from config (pure data, safe at import time).
# MOVE_MAP[direction] = (vx, vy, vyaw). STEP_DURATION_SEC is how long one
# "step" of a bare directional command lasts (2 s at default velocities
# ≈ 60 cm forward or 34° turn). Both live in config_Navigation.json so
# motion can be retuned without editing Python.
_nav = load_config("Navigation")
MOVE_MAP           = {k: tuple(v) for k, v in _nav["move_map"].items()}
STEP_DURATION_SEC  = float(_nav.get("step_duration_sec", 2.0))