JackBot/Robot.py

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

from typing import Protocol, Optional
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
import config as cfg


class RobotBackend(Protocol):
    """Abstraction layer for hardware vs simulation output."""
    def send_angles(self, rad_array: dt.RadArray) -> None:
        ...
    def step_simulation(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


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

    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()


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

    def __init__(
        self,
        backend: Optional[RobotBackend] = None,
        start_pose: str = "init_deg",
        urdf_path: str = cfg.urdf_path
    ):
        # 1. Store configuration & backend FIRST
        self.backend = backend
        self.urdf_path = urdf_path

        # 2. Initialize kinematics and position data
        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
        )

        # 3. Initialize 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]

        # 4. Set state and trigger enter() LAST
        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()


    # -------------------------------------------------------------------------
    # Core Motion Execution
    # -------------------------------------------------------------------------

    def set_joint_angles(self, target_rad: dt.RadArray) -> None:
        """Applies joint angles to internal state and sends to the active backend."""
        self.current_rad = target_rad
        
        if self.backend:
            self.backend.send_angles(target_rad)

    def step_sim(self) -> None:
        """Advances physics simulation step if applicable."""
        if self.backend:
            self.backend.step_simulation()

    def reset_to_init(self) -> None:
        """Resets the robot to its default standing stance."""
        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()

    # -------------------------------------------------------------------------
    # Inverse / Forward Kinematics wrappers bound to this instance
    # -------------------------------------------------------------------------

    def compute_ik(self, target_pos: dt.PosArray) -> dt.RadArray:
        """
        Computes IK using this specific robot instance's current joint angles 
        as the seed initial guess.
        """
        return kin.ikpyInverse(target_pos, initial_rad=self.current_rad)

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

    # -------------------------------------------------------------------------
    # Machine Learning / RL Helper Methods
    # -------------------------------------------------------------------------

    def get_observation(self, command: np.ndarray) -> np.ndarray:
        """
        Returns flat observation vector [joint_angles (18,), command (4,)]
        for reinforcement learning input.
        """
        joint_flat = self.current_rad.data.flatten()
        return np.concatenate([joint_flat, command]).astype(np.float32)

    def apply_rl_action(self, action_delta: np.ndarray, scale: float = math.radians(8.0)) -> dt.RadArray:
        """
        Applies continuous angle deltas from an RL policy network.
        """
        current_flat = self.current_rad.data.flatten()
        new_flat = current_flat + action_delta * scale
        new_rad = dt.RadArray(new_flat.reshape(6, 3))
        self.set_joint_angles(new_rad)
        return new_rad

    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:
        """Executes one step of the current active state."""
        next_state_key = self.current_state.execute(self)
        if next_state_key:
            self.transition_to(next_state_key)
        self.step_sim()