FusionAGI/fusionagi/maa/embodiment.py

"""Embodied Intelligence — robotics bridge for physical actuator integration.

Connects FusionAGI's reasoning and planning pipeline to physical
actuators through a protocol-based abstraction. Supports:
- Robotic arm control (joint positions, trajectories)
- Sensor data ingestion (cameras, LIDAR, IMU)
- Environment perception (object detection, spatial mapping)
- Advisory safety observations (force limits, workspace bounds — logged, not enforced)
"""

from __future__ import annotations

from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from enum import Enum
from typing import Any

from pydantic import BaseModel, Field

from fusionagi._logger import logger


class ActuatorState(str, Enum):
    """Physical actuator operational state."""

    IDLE = "idle"
    MOVING = "moving"
    HOLDING = "holding"
    ERROR = "error"
    EMERGENCY_STOP = "emergency_stop"


class SensorType(str, Enum):
    """Types of physical sensors."""

    CAMERA = "camera"
    LIDAR = "lidar"
    IMU = "imu"
    FORCE_TORQUE = "force_torque"
    PROXIMITY = "proximity"
    TEMPERATURE = "temperature"
    ENCODER = "encoder"


class SensorReading(BaseModel):
    """Single sensor reading with metadata."""

    sensor_id: str = Field(..., description="Unique sensor identifier")
    sensor_type: SensorType = Field(..., description="Type of sensor")
    value: Any = Field(..., description="Sensor value (type depends on sensor)")
    timestamp: float = Field(..., description="Timestamp in seconds")
    confidence: float = Field(default=1.0, ge=0.0, le=1.0, description="Reading confidence")
    metadata: dict[str, Any] = Field(default_factory=dict)


class JointState(BaseModel):
    """State of a single robotic joint."""

    joint_id: str = Field(..., description="Joint identifier")
    position: float = Field(default=0.0, description="Current position (radians or meters)")
    velocity: float = Field(default=0.0, description="Current velocity")
    effort: float = Field(default=0.0, description="Current effort/torque")
    min_limit: float = Field(default=-3.14159, description="Minimum position limit")
    max_limit: float = Field(default=3.14159, description="Maximum position limit")


class TrajectoryPoint(BaseModel):
    """Single point in a motion trajectory."""

    joint_positions: dict[str, float] = Field(default_factory=dict)
    time_from_start: float = Field(default=0.0, description="Seconds from trajectory start")
    velocity: dict[str, float] = Field(default_factory=dict)


class MotionCommand(BaseModel):
    """Command to execute a physical motion."""

    command_id: str = Field(..., description="Unique command identifier")
    trajectory: list[TrajectoryPoint] = Field(default_factory=list)
    max_velocity: float = Field(default=1.0, description="Max velocity scaling [0, 1]")
    max_force: float = Field(default=100.0, description="Max force limit (N)")
    enable_collision_check: bool = Field(default=True)
    metadata: dict[str, Any] = Field(default_factory=dict)


class MotionResult(BaseModel):
    """Result of a motion command execution."""

    command_id: str
    success: bool
    final_joint_states: dict[str, JointState] = Field(default_factory=dict)
    execution_time: float = Field(default=0.0, description="Total execution time (seconds)")
    error_message: str = Field(default="")


class ActuatorAdapter(ABC):
    """Abstract adapter for physical actuator control.

    Implementations connect to specific robots (ROS2, direct serial, etc.).
    """

    @abstractmethod
    async def get_joint_states(self) -> list[JointState]:
        """Read current joint states from hardware."""
        ...

    @abstractmethod
    async def execute_motion(self, command: MotionCommand) -> MotionResult:
        """Execute a motion command on the hardware."""
        ...

    @abstractmethod
    async def emergency_stop(self) -> bool:
        """Trigger emergency stop on all actuators."""
        ...

    @abstractmethod
    async def get_state(self) -> ActuatorState:
        """Get current actuator operational state."""
        ...


class SensorAdapter(ABC):
    """Abstract adapter for sensor data ingestion."""

    @abstractmethod
    async def read(self, sensor_id: str) -> SensorReading | None:
        """Read current value from a sensor."""
        ...

    @abstractmethod
    async def list_sensors(self) -> list[str]:
        """List available sensor IDs."""
        ...


class SimulatedActuator(ActuatorAdapter):
    """Simulated actuator for testing without hardware."""

    def __init__(self, joint_ids: list[str] | None = None) -> None:
        self._joint_ids = joint_ids or ["joint_0", "joint_1", "joint_2", "joint_3"]
        self._states: dict[str, JointState] = {
            jid: JointState(joint_id=jid) for jid in self._joint_ids
        }
        self._actuator_state = ActuatorState.IDLE

    async def get_joint_states(self) -> list[JointState]:
        return list(self._states.values())

    async def execute_motion(self, command: MotionCommand) -> MotionResult:
        self._actuator_state = ActuatorState.MOVING
        for point in command.trajectory:
            for jid, pos in point.joint_positions.items():
                if jid in self._states:
                    state = self._states[jid]
                    clamped = max(state.min_limit, min(state.max_limit, pos))
                    state.position = clamped

        self._actuator_state = ActuatorState.IDLE
        logger.info("Simulated motion executed", extra={"command_id": command.command_id})
        return MotionResult(
            command_id=command.command_id,
            success=True,
            final_joint_states=dict(self._states),
            execution_time=sum(p.time_from_start for p in command.trajectory[-1:]),
        )

    async def emergency_stop(self) -> bool:
        self._actuator_state = ActuatorState.EMERGENCY_STOP
        logger.warning("EMERGENCY STOP triggered (simulated)")
        return True

    async def get_state(self) -> ActuatorState:
        return self._actuator_state


class SimulatedSensor(SensorAdapter):
    """Simulated sensor adapter for testing."""

    def __init__(self) -> None:
        self._sensors: dict[str, SensorReading] = {}

    def register_sensor(self, sensor_id: str, sensor_type: SensorType, value: Any) -> None:
        """Register a simulated sensor."""
        import time

        self._sensors[sensor_id] = SensorReading(
            sensor_id=sensor_id,
            sensor_type=sensor_type,
            value=value,
            timestamp=time.monotonic(),
        )

    async def read(self, sensor_id: str) -> SensorReading | None:
        return self._sensors.get(sensor_id)

    async def list_sensors(self) -> list[str]:
        return list(self._sensors.keys())


@dataclass
class EmbodimentBridge:
    """Bridge between FusionAGI reasoning and physical world.

    Coordinates sensor data ingestion, motion planning integration
    with the MAA pipeline, and actuator command execution with
    safety interlocks.
    """

    actuator: ActuatorAdapter | None = None
    sensors: SensorAdapter | None = None
    workspace_bounds: dict[str, tuple[float, float]] = field(default_factory=dict)
    max_force_limit: float = 150.0
    _command_history: list[MotionResult] = field(default_factory=list)

    async def perceive(self) -> dict[str, Any]:
        """Gather current perception from all sensors and actuator state.

        Returns:
            Dict with sensor readings and joint states.
        """
        perception: dict[str, Any] = {"sensors": {}, "joints": [], "actuator_state": "unknown"}

        if self.actuator:
            perception["actuator_state"] = (await self.actuator.get_state()).value
            perception["joints"] = [j.model_dump() for j in await self.actuator.get_joint_states()]

        if self.sensors:
            sensor_ids = await self.sensors.list_sensors()
            for sid in sensor_ids:
                reading = await self.sensors.read(sid)
                if reading:
                    perception["sensors"][sid] = reading.model_dump()

        return perception

    async def execute(self, command: MotionCommand) -> MotionResult:
        """Execute a motion command with advisory observations.

        Force limits and workspace bounds are logged as advisories
        but do not prevent execution.  The physical hardware has its
        own limits; the software layer observes and learns.

        Args:
            command: Motion command to execute.

        Returns:
            Execution result.
        """
        if not self.actuator:
            return MotionResult(
                command_id=command.command_id,
                success=False,
                error_message="No actuator connected",
            )

        if command.max_force > self.max_force_limit:
            logger.info(
                "Force advisory: commanded force exceeds soft limit (proceeding)",
                extra={
                    "requested": command.max_force,
                    "limit": self.max_force_limit,
                    "mode": "advisory",
                },
            )

        if self.workspace_bounds:
            for point in command.trajectory:
                for jid, pos in point.joint_positions.items():
                    if jid in self.workspace_bounds:
                        lo, hi = self.workspace_bounds[jid]
                        if pos < lo or pos > hi:
                            logger.info(
                                "Workspace advisory: joint outside bounds (proceeding)",
                                extra={
                                    "joint": jid,
                                    "position": pos,
                                    "bounds": [lo, hi],
                                    "mode": "advisory",
                                },
                            )

        result = await self.actuator.execute_motion(command)
        self._command_history.append(result)
        return result

    async def stop(self) -> bool:
        """Emergency stop all actuators."""
        if self.actuator:
            return await self.actuator.emergency_stop()
        return False

    def get_summary(self) -> dict[str, Any]:
        """Return bridge summary."""
        return {
            "actuator_connected": self.actuator is not None,
            "sensors_connected": self.sensors is not None,
            "workspace_bounds": self.workspace_bounds,
            "max_force_limit": self.max_force_limit,
            "commands_executed": len(self._command_history),
        }


__all__ = [
    "ActuatorAdapter",
    "ActuatorState",
    "EmbodimentBridge",
    "JointState",
    "MotionCommand",
    "MotionResult",
    "SensorAdapter",
    "SensorReading",
    "SensorType",
    "SimulatedActuator",
    "SimulatedSensor",
    "TrajectoryPoint",
]