spawn_commands.py

"""JSON helpers for forwarding spawn/reset commands to the world worker."""
from __future__ import annotations

import json
from typing import Any, Optional

from avlite.c10_perception.c11_perception_model import AgentState


def agent_state_to_dict(agent: AgentState) -> dict[str, Any]:
    return {
        "action": "spawn",
        "x": float(agent.x),
        "y": float(agent.y),
        "theta": float(agent.theta),
        "velocity": float(agent.velocity),
        "length": float(agent.length),
        "width": float(agent.width),
        "id": int(agent.agent_id),
    }


def agent_state_from_dict(data: dict[str, Any]) -> AgentState:
    return AgentState(
        x=float(data.get("x", 0)),
        y=float(data.get("y", 0)),
        theta=float(data.get("theta", 0)),
        velocity=float(data.get("velocity", 0)),
        length=float(data.get("length", 4.5)),
        width=float(data.get("width", 2.0)),
        agent_id=int(data.get("id", data.get("agent_id", 0))),
    )


def encode_spawn_command(agent: AgentState) -> str:
    return json.dumps(agent_state_to_dict(agent))


def parse_spawn_command(payload: str) -> Optional[AgentState]:
    data = json.loads(payload)
    if data.get("action", "spawn") != "spawn":
        return None
    return agent_state_from_dict(data)


def encode_teleport_command(x: float, y: float, theta: Optional[float] = None) -> str:
    payload: dict[str, Any] = {
        "action": "teleport",
        "x": float(x),
        "y": float(y),
    }
    if theta is not None:
        payload["theta"] = float(theta)
    return json.dumps(payload)


def parse_teleport_command(payload: str) -> Optional[tuple[float, float, Optional[float]]]:
    data = json.loads(payload)
    if data.get("action") != "teleport":
        return None
    theta = data.get("theta")
    return (
        float(data.get("x", 0)),
        float(data.get("y", 0)),
        float(theta) if theta is not None else None,
    )


def agents_from_perception_json(data: dict[str, Any]) -> list[AgentState]:
    agents = []
    for obj in data.get("objects", []):
        agents.append(agent_state_from_dict(obj))
    return agents