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
import pybullet as p

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, body_id: Optional[int] = None):
        self.sim = sim_instance
        self.body_id = body_id

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

        if self.body_id is not None and hasattr(self.sim, 'updatePosForBody'):
            self.sim.updatePosForBody(self.body_id, rad_array)
        elif hasattr(self.sim, 'updatePos'):
            self.sim.updatePos(rad_array)
        elif hasattr(self.sim, 'set_robot_joint_angles') and self.body_id is not None:
            joint_indices = getattr(self.sim, 'joint_indices', list(range(18)))
            self.sim.set_robot_joint_angles(self.body_id, joint_indices, rad_array.data.flatten())

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

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


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
        self.start_pose = start_pose

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

        # RL configuration
        self.action_scale = 0.1  # Joint delta step size (radians)
        self.gait_phase = 0.0

        # 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:
        pose_deg = ri.init_deg if self.start_pose == "init_deg" else ri.init90_deg
        self.current_rad = pose_deg.to_rad()
        self.current_pos = kin.ikpyForward(self.current_rad)
        self.gait_phase = 0.0
        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 step_kinematic_gait(self, vx: float, vy: float, omega: float) -> None:
        """Procedural Tripod Gait IK solver for kinematic execution & evaluation."""
        cmd_mag = math.hypot(vx, vy) + abs(omega)
        if cmd_mag < 0.03:
            target_rad = self.compute_ik(self.center_points)
            self.set_joint_angles(target_rad)
            return

        # Advance gait step cycle phase
        self.gait_phase = (self.gait_phase + 0.22) % (2.0 * math.pi)

        stride_len = 0.045   # 4.5 cm maximum stride
        step_height = 0.035  # 3.5 cm foot clearance height

        center_data = (
            self.center_points.data 
            if hasattr(self.center_points, 'data') 
            else np.array(self.center_points)
        )

        target_positions = []
        for leg_id in range(6):
            base_pos = np.array(center_data[leg_id], dtype=np.float32)
            
            # Tripod leg grouping phase offset (even vs odd leg IDs)
            phase_offset = 0.0 if (leg_id % 2 == 0) else math.pi
            leg_phase = (self.gait_phase + phase_offset) % (2.0 * math.pi)

            # Directional motion unit vector calculation
            lx, ly = base_pos[0], base_pos[1]
            rot_dx = -omega * ly
            rot_dy = omega * lx

            dx_dir = vx + rot_dx
            dy_dir = vy + rot_dy
            dir_norm = math.hypot(dx_dir, dy_dir) + 1e-6

            dx_unit = dx_dir / dir_norm
            dy_unit = dy_dir / dir_norm

            if leg_phase < math.pi:
                # Swing Phase (Leg lifted & stepping forward)
                progress = math.cos(leg_phase)
                lift = math.sin(leg_phase) * step_height
                dx = -progress * stride_len * dx_unit
                dy = -progress * stride_len * dy_unit
                dz = lift
            else:
                # Stance Phase (Leg grounded & propelling torso)
                progress = math.cos(leg_phase - math.pi)
                dx = progress * stride_len * dx_unit
                dy = progress * stride_len * dy_unit
                dz = 0.0

            target_leg_pos = base_pos + np.array([dx, dy, dz], dtype=np.float32)
            target_positions.append(target_leg_pos)

        target_pos_array = dt.PosArray(np.array(target_positions))
        target_rad = self.compute_ik(target_pos_array)
        self.set_joint_angles(target_rad)

    def step_with_command(
        self, 
        command: np.ndarray, 
        action: Optional[np.ndarray] = None, 
        mode: str = "direct"
    ) -> None:
        """
        Unified motion execution method supporting Direct RL, Residual RL, and Pure Kinematics.

        Args:
            command: np.ndarray [vx, vy, vz, omega] from RL environment
            action: np.ndarray [18,] RL action deltas from neural network ([-1, 1])
            mode: "kinematics_only" | "residual" | "direct"
        """
        # 1. Map RL command [vx, vy, vz, omega] to Robot motion vector [vx, vy, omega]
        cmd_vx, cmd_vy, _, cmd_omega = command
        self.vector_dirmov = [float(cmd_vx), float(cmd_vy), float(cmd_omega)]

        # 2. Automatically trigger state transitions based on command magnitude
        cmd_magnitude = math.hypot(cmd_vx, cmd_vy) + abs(cmd_omega)
        if cmd_magnitude > 0.05 and self.current_state_key == "idle":
            # Transition to walking state if registered in state machine
            target_state = "walk" if "walk" in STATE_REGISTRY else "move"
            if target_state in STATE_REGISTRY:
                self.transition_to(target_state)
        elif cmd_magnitude <= 0.05 and self.current_state_key != "idle":
            self.transition_to("idle")

        # 3. Execute according to chosen mode
        if mode == "kinematics_only":
            self.step_kinematic_gait(cmd_vx, cmd_vy, cmd_omega)

        elif mode == "residual":
            if self.current_state_key != "idle" and self.current_state and self.current_state_key in STATE_REGISTRY:
                self.current_state.execute(self)
            else:
                self.step_kinematic_gait(cmd_vx, cmd_vy, cmd_omega)

            if action is not None:
                action_flat = np.clip(np.asarray(action, dtype=np.float32), -1.0, 1.0) * self.action_scale
                kin_flat = self.current_rad.data.flatten()
                final_flat = np.clip(kin_flat + action_flat, -np.pi / 2, np.pi / 2)
                self.set_joint_angles(dt.RadArray(data=final_flat.reshape(self.current_rad.data.shape)))

        elif mode == "direct":
            if action is not None:
                self.apply_rl_action(action)

    # --- RL METHODS ---

    def apply_rl_action(self, action: np.ndarray) -> None:
        action = np.asarray(action, dtype=np.float32)
        scaled_action = np.clip(action, -1.0, 1.0) * self.action_scale

        # Sync with actual PyBullet state if backend supports it
        if isinstance(self.backend, PyBulletBackend) and self.backend.sim:
            actual_angles = self.backend.sim.get_robot_joint_angles(self.backend.body_id)
            current_flat = actual_angles.flatten()
        else:
            current_flat = self.current_rad.data.flatten()

        updated_flat = np.clip(current_flat + scaled_action, -np.pi / 2, np.pi / 2)
        new_rad = dt.RadArray(data=updated_flat.reshape(self.current_rad.data.shape))
        self.set_joint_angles(new_rad)

    def get_observation(self, command: Optional[np.ndarray] = None) -> np.ndarray:
        """
        Returns observation vector [18 joint angles] + [optional 4 command dimensions].
        Queries SimManager helper if PyBulletBackend is used; falls back to internal state otherwise.
        """
        if isinstance(self.backend, PyBulletBackend) and self.backend.sim:
            body_id = self.backend.body_id if self.backend.body_id is not None else 0
            if hasattr(self.backend.sim, 'get_robot_joint_angles'):
                joint_angles = self.backend.sim.get_robot_joint_angles(body_id)
            elif hasattr(self.backend.sim, 'physics_client'):
                physics_client = self.backend.sim.physics_client
                joint_indices = self.backend.sim.robot_joints.get(body_id, list(range(18))) if hasattr(self.backend.sim, 'robot_joints') else list(range(18))
                joint_states = p.getJointStates(body_id, joint_indices, physicsClientId=physics_client)
                joint_angles = np.array([state[0] for state in joint_states], dtype=np.float32)
            else:
                joint_angles = self.current_rad.data.flatten().astype(np.float32)
        else:
            joint_angles = self.current_rad.data.flatten().astype(np.float32)

        if command is not None:
            cmd = np.asarray(command, dtype=np.float32).flatten()
            return np.concatenate([joint_angles, cmd]).astype(np.float32)

        return joint_angles.astype(np.float32)

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