p48_node_bootstrap.py

"""Bootstrap AVLite stack components for ROS node worker processes."""
from __future__ import annotations

import inspect
import logging
import os
from dataclasses import dataclass
from typing import Optional

from avlite.c10_perception.c11_perception_model import EgoState, PerceptionModel, Map, RaceMap
from avlite.c10_perception.c12_perception_strategy import PerceptionStrategy
from avlite.c20_planning.c21_planning_model import GlobalPlan
from avlite.c20_planning.c22_global_planning_strategy import GlobalPlannerStrategy
from avlite.c20_planning.c23_local_planning_strategy import LocalPlanningStrategy
from avlite.c20_planning.c24_global_hdmap_planners import HDMapGlobalPlanner
from avlite.c20_planning.c25_global_race_planners import GlobalCenterlineRacePlanner
from avlite.c20_planning.c27_local_behavioral_and_velocity_planners import VelocityLocalPlanner  # noqa: F401 — registry
from avlite.c20_planning.c28_local_lattice_planners import GreedyLatticePlanner
from avlite.c30_control.c32_control_strategy import ControlStrategy
from avlite.c30_control.c33_pid import PIDController  # noqa: F401 — registry
from avlite.c30_control.c34_stanley import StanleyController
from avlite.c10_perception.c15_perception_algs import ConstantVelocityPrediction  # noqa: F401 — registry
from avlite.c40_execution.c41_world_bridge import WorldBridge
from avlite.c60_apps.c62_factory import load_stack_settings
from avlite.c40_execution.c46_basic_sim import BasicSim
from avlite.c40_execution.c49_settings import ExecutionSettings
from avlite.c60_apps.c68_paths import DataPaths
from avlite.c60_apps.c65_setting_utils import load_setting

from .settings import PluginSettings

log = logging.getLogger(__name__)

VALID_ROLES = frozenset({"world", "planner", "controller", "perception"})
_LEGACY_PROXY_PLANNERS = frozenset({"ROSLocalPlanner", "ProxyLocalPlanner"})
_LEGACY_PROXY_CONTROLLERS = frozenset({"ROSController", "ProxyController"})


def _worker_planner_name() -> str:
    """Resolve the compute planner class for the planner worker process."""
    env_name = os.environ.get("AVLITE_LOCAL_PLANNER")
    if env_name:
        return env_name
    if PluginSettings.compute_planner:
        return PluginSettings.compute_planner
    name = ExecutionSettings.c40_local_planner
    if name in _LEGACY_PROXY_PLANNERS:
        return GreedyLatticePlanner.__name__
    return name


def _worker_controller_name() -> str:
    """Resolve the compute controller class for the controller worker process."""
    env_name = os.environ.get("AVLITE_CONTROLLER")
    if env_name:
        return env_name
    if PluginSettings.compute_controller:
        return PluginSettings.compute_controller
    name = ExecutionSettings.c40_controller
    if name in _LEGACY_PROXY_CONTROLLERS:
        return StanleyController.__name__
    return name


def _instantiate_planner(cls, global_plan: GlobalPlan, pm: PerceptionModel):
    params = inspect.signature(cls.__init__).parameters
    if "env" in params:
        return cls(global_plan=global_plan, env=pm)
    return cls(global_plan=global_plan, pm=pm)


def _bootstrap_perception(pm: PerceptionModel) -> Optional[PerceptionStrategy]:
    name = ExecutionSettings.c40_perception
    if name and name in PerceptionStrategy.registry:
        cls = PerceptionStrategy.registry[name]
        return cls(perception_model=pm)
    return None


@dataclass
class NodeBootstrapResult:
    ego_state: EgoState
    pm: PerceptionModel
    world: Optional[WorldBridge] = None
    local_planner: Optional[LocalPlanningStrategy] = None
    controller: Optional[ControlStrategy] = None
    perception: Optional[PerceptionStrategy] = None
    global_plan: Optional[GlobalPlan] = None


def profile_from_env() -> str:
    return os.environ.get("AVLITE_PROFILE", "default")


def ego_state_from_env(ego_state: EgoState) -> EgoState:
    """Override bootstrap ego from AVLITE_EGO_* env (set by main process on worker start)."""
    if "AVLITE_EGO_X" not in os.environ:
        return ego_state
    ego_state.x = float(os.environ["AVLITE_EGO_X"])
    ego_state.y = float(os.environ["AVLITE_EGO_Y"])
    if "AVLITE_EGO_THETA" in os.environ:
        ego_state.theta = float(os.environ["AVLITE_EGO_THETA"])
    return ego_state


def _bridge_kwargs(bridge_cls, ego_state: EgoState, world_pm: PerceptionModel, loaded_map) -> dict:
    def _reference_point_tuple():
        ref = ExecutionSettings.c40_reference_point
        if ref and len(ref) >= 2:
            return float(ref[0]), float(ref[1])
        return None

    params = inspect.signature(bridge_cls.__init__).parameters
    kwargs: dict = {}
    if "ego_state" in params:
        kwargs["ego_state"] = ego_state
    if "pm" in params:
        kwargs["pm"] = world_pm
    if "reference_point" in params:
        kwargs["reference_point"] = _reference_point_tuple()
    if "map" in params:
        kwargs["map"] = loaded_map
    return kwargs


def _load_map():
    map_file = os.environ.get("AVLITE_MAP") or ExecutionSettings.c40_map
    if not map_file:
        return None
    if os.environ.get("AVLITE_MAP"):
        log.info("Loading map from AVLITE_MAP=%s", map_file)
    return Map.open(DataPaths.resolve_stored(map_file))


def _load_global_plan(loaded_map) -> GlobalPlan:
    plan_file = os.environ.get("AVLITE_GLOBAL_TRAJECTORY") or ExecutionSettings.c40_global_trajectory
    if not plan_file:
        log.warning("No global trajectory configured; using race centerline fallback if available.")
        if isinstance(loaded_map, RaceMap):
            gp = GlobalCenterlineRacePlanner(map=loaded_map)
            return gp.plan()
        return GlobalPlan()
    try:
        path = DataPaths.resolve_stored(plan_file)
        if os.environ.get("AVLITE_GLOBAL_TRAJECTORY"):
            log.info("Loading global plan from AVLITE_GLOBAL_TRAJECTORY=%s", plan_file)
        return GlobalPlan.from_file(path)
    except Exception as e:
        log.warning("Could not load global trajectory %s: %s. Using race centerline.", plan_file, e)
        if isinstance(loaded_map, RaceMap):
            gp = GlobalCenterlineRacePlanner(map=loaded_map)
            return gp.plan()
        log.warning("No RaceMap available for centerline fallback; using empty GlobalPlan.")
        return GlobalPlan()


def bootstrap_role(role: str, profile: Optional[str] = None) -> NodeBootstrapResult:
    """Load settings and construct stack objects for a ROS worker role."""
    if role not in VALID_ROLES:
        raise ValueError(f"Unknown role {role!r}; expected one of {sorted(VALID_ROLES)}")

    profile = profile or profile_from_env()
    load_stack_settings(profile=profile, load_plugins=True)
    load_setting(PluginSettings, profile=profile)

    loaded_map = _load_map()
    default_global_plan = _load_global_plan(loaded_map)
    ego_state = EgoState(x=default_global_plan.start_point[0], y=default_global_plan.start_point[1])
    ego_state = ego_state_from_env(ego_state)
    pm = PerceptionModel(ego_vehicle=ego_state)
    if loaded_map is not None:
        pm.map = loaded_map

    result = NodeBootstrapResult(
        ego_state=ego_state,
        pm=pm,
        global_plan=default_global_plan,
    )

    if role == "world":
        world_pm = PerceptionModel(ego_vehicle=ego_state)
        if loaded_map is not None:
            world_pm.map = loaded_map
        bridge = ExecutionSettings.c40_bridge
        if bridge in WorldBridge.registry:
            cls = WorldBridge.registry[bridge]
            result.world = cls(**_bridge_kwargs(cls, ego_state, world_pm, loaded_map))
        else:
            result.world = BasicSim(**_bridge_kwargs(BasicSim, ego_state, world_pm, loaded_map))
        return result

    if role == "planner":
        local_global_plan = default_global_plan
        if ExecutionSettings.c40_global_planner == HDMapGlobalPlanner.__name__:
            local_global_plan = GlobalPlan(
                start_point=default_global_plan.start_point,
                goal_point=default_global_plan.goal_point,
                path=default_global_plan.path,
                velocity=default_global_plan.velocity,
                trajectory=default_global_plan.trajectory,
            )
        name = _worker_planner_name()
        log.info("Bootstrapping planner worker with %s", name)
        try:
            if name in LocalPlanningStrategy.registry:
                cls = LocalPlanningStrategy.registry[name]
                result.local_planner = _instantiate_planner(cls, local_global_plan, pm)
            else:
                log.warning("Planner %s not in registry; using GreedyLatticePlanner", name)
                result.local_planner = GreedyLatticePlanner(global_plan=local_global_plan, env=pm)
        except Exception as e:
            log.error("Planner bootstrap failed for %s: %s", name, e)
            result.local_planner = GreedyLatticePlanner(global_plan=local_global_plan, env=pm)
        result.perception = _bootstrap_perception(pm)
        if result.perception is not None:
            log.info("Planner worker perception strategy: %s", ExecutionSettings.c40_perception)
        return result

    if role == "controller":
        name = _worker_controller_name()
        log.info("Bootstrapping controller worker with %s", name)
        try:
            if name in ControlStrategy.registry:
                cls = ControlStrategy.registry[name]
                result.controller = cls()
            else:
                log.warning("Controller %s not in registry; using StanleyController", name)
                result.controller = StanleyController()
            if default_global_plan.trajectory is not None:
                result.controller.set_trajectory_tracker(default_global_plan.trajectory)
        except Exception as e:
            log.error("Controller bootstrap failed for %s: %s", name, e)
            result.controller = StanleyController()
            if default_global_plan.trajectory is not None:
                result.controller.set_trajectory_tracker(default_global_plan.trajectory)
        return result

    if role == "perception":
        name = ExecutionSettings.c40_perception
        if name and name in PerceptionStrategy.registry:
            cls = PerceptionStrategy.registry[name]
            result.perception = cls(perception_model=pm)
        return result

    return result


def attach_worker_logging(node, logger_prefix: str) -> None:
    """Forward avlite module logs to rosout so the UI collector can display them."""
    del logger_prefix  # all avlite.* loggers share one forward handler
    class _ForwardFilter(logging.Filter):
        """Drop high-rate tick logs that would flood rosout and the spin thread."""
        _SKIP_PREFIXES = ("Step:",)

        def filter(self, record: logging.LogRecord) -> bool:
            msg = record.getMessage()
            return not any(msg.startswith(prefix) for prefix in self._SKIP_PREFIXES)

    class _RosForwardHandler(logging.Handler):
        avlite_ros_forward = True

        def __init__(self, ros_logger):
            super().__init__()
            self._ros = ros_logger
            self.addFilter(_ForwardFilter())

        def emit(self, record: logging.LogRecord) -> None:
            try:
                msg = record.getMessage()
                if record.levelno >= logging.ERROR:
                    self._ros.error(msg)
                elif record.levelno >= logging.WARNING:
                    self._ros.warning(msg)
                else:
                    self._ros.info(msg)
            except Exception:
                self.handleError(record)

    target = logging.getLogger("avlite")
    if any(getattr(h, "avlite_ros_forward", False) for h in target.handlers):
        return
    target.addHandler(_RosForwardHandler(node.get_logger()))
    target.setLevel(logging.INFO)


def spin_node(node) -> None:
    """Run rclpy.spin until shutdown (Ctrl+C, SIGTERM, or rclpy.shutdown())."""
    import rclpy
    from rclpy.executors import ExternalShutdownException

    try:
        rclpy.spin(node)
    except (KeyboardInterrupt, ExternalShutdownException):
        pass