JackBot/Robot.py

"""
Robot.py - Unified Robot Class for JackBot
Handles state, kinematics, backends (Hardware/Simulation), and motion execution.
"""

from typing import Protocol, Optional, Union
import numpy as np
import math

from states import STATE_REGISTRY
from states.State import State
import DataTypes as dt
import kinematics as kin
import robot_init as ri
from config import cfg, BackendType

# Import communications and simulation modules
from simulation import Simulation
from EspCommunication import ESP32Communication
from ArduinoCommunication import ArduinoCommunication


class RobotBackend(Protocol):
    """Abstraction layer for hardware vs simulation output."""
    def send_angles(self, rad_array: dt.RadArray) -> None:
        ...
    def step_simulation(self) -> None:
        ...
    def cleanup(self) -> None:
        ...


class HardwareBackend:
    """Backend for physical ESP32 or Arduino robot."""
    def __init__(self, comm_channel):
        self.comm_channel = comm_channel

    def send_angles(self, rad_array: dt.RadArray) -> None:
        if self.comm_channel:
            self.comm_channel.send_motion(rad_array)

    def step_simulation(self) -> None:
        pass  # Physical hardware steps in real-time

    def cleanup(self) -> None:
        if self.comm_channel and hasattr(self.comm_channel, 'close'):
            self.comm_channel.close()


class PyBulletBackend:
    """Backend for PyBullet simulation execution."""
    def __init__(self, sim_instance):
        self.sim = sim_instance

    def send_angles(self, rad_array: dt.RadArray) -> None:
        if self.sim:
            self.sim.updatePos(rad_array)

    def step_simulation(self) -> None:
        if self.sim:
            self.sim.step()

    def cleanup(self) -> None:
        if self.sim and hasattr(self.sim, 'close'):
            self.sim.close()


class Robot:
    """
    Encapsulates a single JackBot hexapod instance.
    Maintains joint states, leg positions, kinematics, and backend control.
    """

    def __init__(
        self,
        backend_type: Union[BackendType, RobotBackend] = BackendType.SIMULATION,
        start_pose: str = "init_deg",
        urdf_path: str = cfg.urdf_path
    ):
        self.urdf_path = urdf_path

        # --- BACKEND FACTORY CREATION ---
        if isinstance(backend_type, BackendType):
            if backend_type == BackendType.SIMULATION:
                # Launch PyBullet 3D Simulation GUI
                sim_instance = Simulation(urdf_path=self.urdf_path)
                self.backend: RobotBackend = PyBulletBackend(sim_instance)
            elif backend_type == BackendType.ESP32:
                comm = ESP32Communication(ip=cfg.esp32_ip, port=cfg.esp32_port)
                self.backend = HardwareBackend(comm)
            elif backend_type == BackendType.ARDUINO:
                comm = ArduinoCommunication(port=cfg.port, baudrate=cfg.baudrate)
                self.backend = HardwareBackend(comm)
        else:
            self.backend = backend_type

        # Kinematics and position initialization
        pose_deg = ri.init_deg if start_pose == "init_deg" else ri.init90_deg
        self.current_rad: dt.RadArray = pose_deg.to_rad()
        self.current_pos: dt.PosArray = kin.ikpyForward(self.current_rad)
        self.center_points: dt.PosArray = (
            ri.get_center_points() if hasattr(ri, "get_center_points") else ri.center_points
        )

        # Gait / motion variables
        self.leg_state = np.array(["step", "drag", "step", "drag", "step", "drag"])
        self.robot_state = "idle"
        self.vector_dirmov = [0.0, 0.0, 0.0]  # [vx, vy, omega]

        # State machine initialization
        self.current_state_key: str = "idle"
        self.current_state: State = STATE_REGISTRY["idle"]
        self.current_state.enter(self)

    def change_state(self, new_state: State) -> None:
        if self.current_state:
            self.current_state.exit(self)
        self.current_state = new_state
        self.current_state.enter(self)

    def update(self) -> None:
        if self.current_state:
            self.current_state.execute(self)
        self.step_sim()

    def set_joint_angles(self, target_rad: dt.RadArray) -> None:
        self.current_rad = target_rad
        if self.backend:
            self.backend.send_angles(target_rad)

    def step_sim(self) -> None:
        if self.backend:
            self.backend.step_simulation()

    def reset_to_init(self) -> None:
        self.current_rad = ri.init_deg.to_rad()
        self.current_pos = kin.ikpyForward(self.current_rad)
        self.set_joint_angles(self.current_rad)
        self.step_sim()

    def compute_ik(self, target_pos: dt.PosArray) -> dt.RadArray:
        return kin.ikpyInverse(target_pos, initial_rad=self.current_rad)

    def compute_fk(self, target_rad: Optional[dt.RadArray] = None) -> dt.PosArray:
        rads = target_rad if target_rad is not None else self.current_rad
        return kin.ikpyForward(rads)

    def transition_to(self, next_state_key: str) -> None:
        if next_state_key in STATE_REGISTRY and next_state_key != self.current_state_key:
            self.current_state.exit(self)
            self.current_state_key = next_state_key
            self.current_state = STATE_REGISTRY[next_state_key]
            self.current_state.enter(self)

    def tick(self) -> None:
        next_state_key = self.current_state.execute(self)
        if next_state_key:
            self.transition_to(next_state_key)
        self.step_sim()

    def cleanup(self) -> None:
        if self.backend and hasattr(self.backend, 'cleanup'):
            self.backend.cleanup()