topic_mirror_bridge.py

"""World bridge that mirrors state from ROS topics (sim runs in a child process)."""
from __future__ import annotations

from dataclasses import dataclass, field
from typing import Callable, Optional

import numpy as np

from avlite import (
    AgentState,
    ControlCommand,
    EgoState,
    ExecutionSettings,
    GlobalPlan,
    LidarCloud,
    PerceptionModel,
    StackCapability,
    WorldBridge,
    WorldCapability,
)
from avlite.c10_perception.c11_perception_model import Map
from avlite.c40_execution.c46_basic_sim import boundary_segments_from_map
from avlite.c60_apps.c68_paths import DataPaths

SpawnForwarder = Callable[[AgentState], None]
ResetForwarder = Callable[[], None]
TeleportForwarder = Callable[[float, float, Optional[float]], None]


def _caps_from_selected_bridge() -> tuple[frozenset, frozenset, frozenset] | None:
    """Class-level caps from ``ExecutionSettings.c40_bridge`` when plain frozensets."""
    name = ExecutionSettings.c40_bridge
    cls = WorldBridge.registry.get(name) if name else None
    if cls is None:
        return None
    world = getattr(cls, "world_capabilities", None)
    stack = getattr(cls, "stack_capabilities", None)
    reqs = getattr(cls, "stack_requirements", frozenset())
    if isinstance(world, property) or isinstance(stack, property):
        return None
    return (
        frozenset(world or ()),
        frozenset(stack or ()),
        frozenset(reqs or ()),
    )


@dataclass
class TopicMirrorBridge(WorldBridge, abstract=True):
    """Mirror ego/perception in the main process while the sim runs in a ROS worker.

    Capability contracts are copied from the selected (factory-built) world bridge
    so the UI and ``available_stack_capabilities()`` see the same GT surface as
    the worker-side simulator.
    """

    owns_ros_topics: bool = True
    spawn_forwarder: Optional[SpawnForwarder] = field(default=None, repr=False)
    reset_forwarder: Optional[ResetForwarder] = field(default=None, repr=False)
    teleport_forwarder: Optional[TeleportForwarder] = field(default=None, repr=False)
    _lidar_buffer: Optional[np.ndarray] = field(default=None, repr=False)
    mirrored_world_capabilities: frozenset = field(default_factory=frozenset)
    mirrored_stack_capabilities: frozenset = field(default_factory=frozenset)
    mirrored_stack_requirements: frozenset = field(default_factory=frozenset)
    boundary_segments: np.ndarray = field(init=False, repr=False)

    def __post_init__(self) -> None:
        self.refresh_boundary_segments()
        if not self.mirrored_world_capabilities and not self.mirrored_stack_capabilities:
            fallback = _caps_from_selected_bridge()
            if fallback is not None:
                self.mirrored_world_capabilities, self.mirrored_stack_capabilities, self.mirrored_stack_requirements = fallback

    def refresh_boundary_segments(self) -> None:
        """Reload track boundary segments from the unified map setting."""
        path = ExecutionSettings.c40_map
        loaded = Map.open(DataPaths.resolve_stored(path)) if path else None
        self.boundary_segments = boundary_segments_from_map(loaded)

    @property
    def world_capabilities(self) -> set[WorldCapability]:
        return set(self.mirrored_world_capabilities)

    @property
    def stack_capabilities(self) -> set[StackCapability]:
        return set(self.mirrored_stack_capabilities)

    @property
    def stack_requirements(self) -> set:
        return set(self.mirrored_stack_requirements)

    def control_ego_state(self, cmd: ControlCommand, dt: Optional[float] = 0.01):
        """Control is applied by the controller/world ROS worker processes."""
        pass

    def teleport_ego(self, x: float, y: float, theta: Optional[float] = None):
        self.ego_state.x = x
        self.ego_state.y = y
        if theta is not None:
            self.ego_state.theta = theta
        if self.teleport_forwarder is not None:
            self.teleport_forwarder(x, y, theta)

    def get_ground_truth_perception_model(self) -> PerceptionModel:
        if self.perception_model is None:
            self.perception_model = PerceptionModel(ego_vehicle=self.ego_state)
        return self.perception_model

    def get_lidar_data(self) -> LidarCloud | None:
        return self._lidar_buffer

    def spawn_agent(self, agent_state: AgentState, global_plan: Optional[GlobalPlan] = None):
        pm = self.get_ground_truth_perception_model()
        pm.add_agent_vehicle(agent_state)
        if self.spawn_forwarder is not None:
            self.spawn_forwarder(agent_state)

    def reset(self):
        pm = self.get_ground_truth_perception_model()
        pm.reset()
        if self.reset_forwarder is not None:
            self.reset_forwarder()