space_robot.py

import numpy as np
import matplotlib.pyplot as plt
from scipy.spatial.transform import Rotation as R
from scipy.spatial.transform import Slerp
import os

from space_object_abstract import SpaceObjectAbstract
from satellite import Satellite
from robot_arm_controller import RobotArm

class SpaceRobot(SpaceObjectAbstract):
    def __init__(self, env_runner, robot_name_in_sim, target_handle=None):
        super().__init__(env_runner, robot_name_in_sim, target_handle)
        
        self.MAX_FORCE = 500.0     
        self.MAX_TORQUE = 100.0   
        self.PLAN_LIN_VEL = 0.2   
        self.PLAN_ANG_VEL = 0.05  

        self.traj_step_size = 1.0

        self.arm1 = RobotArm(env_runner, 1)
        self.arm2 = RobotArm(env_runner, 2)

        # Temporary storage for the current control input to be logged
        self.current_force = np.zeros(3)
        self.current_torque = np.zeros(3) 
        
    def apply_control(self, force_world, torque_world):
        f_norm = np.linalg.norm(force_world)
        t_norm = np.linalg.norm(torque_world)
        if f_norm > self.MAX_FORCE:
            f_clamped = force_world * (self.MAX_FORCE / f_norm)
        else:
            f_clamped = force_world
        if t_norm > self.MAX_TORQUE:
            t_clamped = torque_world * (self.MAX_TORQUE / t_norm)
        else:
            t_clamped = torque_world

        # Store for logging
        self.current_force = f_clamped
        self.current_torque = t_clamped

        self.env_runner.sim.addForceAndTorque(self.base_handle, f_clamped.tolist(), t_clamped.tolist())

    def step_trajectory_control(self, ref_pos, ref_quat, ref_vel=np.zeros(3), ref_ang_vel=np.zeros(3)):
        curr_pos, curr_vel_w, curr_quat, curr_ang_vel_w = self.get_state()
        self._log_data()
        r_curr = R.from_quat(curr_quat)

        ref_vel = r_curr.inv().apply(ref_vel)
        ref_ang_vel = r_curr.inv().apply(ref_ang_vel)

        r_target = R.from_quat(ref_quat)
        omega_b = r_curr.inv().apply(curr_ang_vel_w)
        rot_err_vec = (r_curr.inv() * r_target).as_rotvec()

        wn_att = 1.0; zeta = 1.0
        Kp_mat = (wn_att**2) * self.inertia_matrix
        Kd_mat = (2 * zeta * wn_att) * self.inertia_matrix
        torque_b = Kp_mat @ rot_err_vec + Kd_mat @ (ref_ang_vel - omega_b) + np.cross(omega_b, self.inertia_matrix @ omega_b)
        torque_w = r_curr.apply(torque_b)

        pos_err = ref_pos - curr_pos
        wn_pos = 1.0; zeta_pos = 1.0
        Kp_pos = wn_pos**2 * self.mass
        Kd_pos = 2 * zeta_pos * wn_pos * self.mass
        force_w = Kp_pos * pos_err + Kd_pos * (ref_vel - curr_vel_w)
        
        self.apply_control(force_w, torque_w)
        return np.linalg.norm(pos_err), np.linalg.norm(rot_err_vec)

    def run_trajectory_sequence(self, waypoints):
        print("\n🔵 [阻塞模式] 启动轨迹规划任务...")
        curr_pos, _, curr_quat, _ = self.get_state()
        for i, (t_pos, t_quat) in enumerate(waypoints):
            print(f"\n📍 航点 {i+1}: {t_pos} ...")
            self._execute_single_waypoint(curr_pos, curr_quat, np.array(t_pos), t_quat)
            curr_pos, _, curr_quat, _ = self.get_state()

    def _execute_single_waypoint(self, start_pos, start_quat, end_pos, end_quat):
        dist = np.linalg.norm(end_pos - start_pos)
        angle_diff = np.linalg.norm((R.from_quat(end_quat) * R.from_quat(start_quat).inv()).as_rotvec())
        duration = max(dist/self.PLAN_LIN_VEL, angle_diff/self.PLAN_ANG_VEL, 5.0)
        
        slerp = Slerp([0, duration], R.from_quat([start_quat, end_quat]))
        start_sim_time = self.env_runner.sim.getSimulationTime()
        while True:
            t = self.env_runner.sim.getSimulationTime() - start_sim_time
            if t > duration:
                curr_ref_pos, curr_ref_quat, finished = end_pos, end_quat, True
            else:
                alpha = t / duration
                curr_ref_pos = start_pos + (end_pos - start_pos) * alpha
                curr_ref_quat = slerp([t])[0].as_quat()
                finished = False
            
            p_err, r_err = self.step_trajectory_control(curr_ref_pos, curr_ref_quat)
            self.env_runner.stepSimulation()
            
            if finished and p_err < 0.02 and r_err < 0.2:
                print("      ✅ 到达")
                break
            if t > duration + 20.0:
                print("      ⚠️ 超时")
                break

    def approach_and_grasp_satellite(self, satellite: Satellite, approach_dist=10.0):
        print(f"\n🟢 [阻塞模式] 靠近目标卫星并同步...")
        grasp_success = False
        to_grasp = False
        steady_timesteps = 0
        while not grasp_success:
            sat_pos, sat_vel, sat_quat, sat_ang_vel = satellite.get_state()
            satellite_rotmat = R.from_quat(sat_quat).as_matrix()
            target_pos = sat_pos - approach_dist * satellite_rotmat[:, 1]
            target_quat = sat_quat

            current_pos, _, current_quat, _ = self.get_state()
            direction_vec = target_pos - current_pos
            distance = np.linalg.norm(direction_vec)
            if distance > self.traj_step_size:
                target_pos = current_pos + (direction_vec / distance) * self.traj_step_size
                alpha = self.traj_step_size / distance
                slerp = Slerp([0, 1], R.from_quat([current_quat, target_quat]))
                target_quat = slerp([alpha])[0].as_quat()

            pos_error, rot_error = self.step_trajectory_control(target_pos, target_quat, sat_vel, sat_ang_vel)

            if not to_grasp and pos_error < 0.5 and rot_error < 0.1:
                steady_timesteps += 1
                if steady_timesteps >= self.env_runner.sim_rate * 10:  # 稳定3秒
                    to_grasp = True
                    print("      🔄 靠近完成，准备抓取...")
            else:
                steady_timesteps = 0

            if to_grasp:
                target_grasp_pos, target_grasp_quat = self.env_runner.getObjectPose("/satellite", self.arm2.base_name)
                target_grasp_rotmat = R.from_quat(target_grasp_quat).as_matrix()
                target_grasp_pos -= 0.1 * target_grasp_rotmat[:, 1]
                target_grasp_pos += 0.4 * target_grasp_rotmat[:, 2]
                grasp_success = self.arm2.step_grasp((target_grasp_pos, target_grasp_quat))

            self.env_runner.stepSimulation()

        self.env_runner.sim.setObjectParent(satellite.base_handle, self.arm2.force_sensor, True)
        self.env_runner.sim.resetDynamicObject(satellite.base_handle) 
        print("      ✅ 抓取完成，已与目标卫星对接")

    def stop_spinning(self):
        print(f"\n🟢 [阻塞模式] 停止目标卫星自旋...")
        robot_pos, robot_lin, robot_quat, robot_ang = self.get_state()
        n_steps = 1000
        for step in range(1, n_steps + 1):
            alpha = step / n_steps
            target_vel = robot_lin * (1 - alpha)
            target_ang_vel = robot_ang * (1 - alpha)
            self.step_trajectory_control(robot_pos, robot_quat, target_vel, target_ang_vel)
            self.env_runner.stepSimulation()

            robot_pos, robot_lin, robot_quat, robot_ang = self.get_state()
            ang_vel_norm = np.linalg.norm(robot_ang)
            if ang_vel_norm < 0.01:
                break
        print("      ✅ 自旋已停止")

    def repair_satellite(self, satellite: Satellite):
        print(f"\n🟢 [阻塞模式] 进行卫星维修操作...")
        
        target_grasp_pos, target_grasp_quat = self.env_runner.getObjectPose("/satellite", self.arm1.base_name)
        target_grasp_rotmat = R.from_quat(target_grasp_quat).as_matrix()
        target_grasp_pos -= 0.4 * target_grasp_rotmat[:, 1]

        target_grasp_pos -= 0.8 * target_grasp_rotmat[:, 2]
        grasp_success = self.arm1.move_to_target_pose((target_grasp_pos, target_grasp_quat), log_plot=True)
        target_grasp_pos += 1.2 * target_grasp_rotmat[:, 0]
        grasp_success = self.arm1.move_to_target_pose((target_grasp_pos, target_grasp_quat))
        target_grasp_pos += 0.25 * target_grasp_rotmat[:, 1]
        grasp_success = self.arm1.move_to_target_pose((target_grasp_pos, target_grasp_quat), duration_s=0.5, settle_s=0.0)
        satellite.openPanel2(max_vel=2.0)
        target_grasp_pos -= 0.25 * target_grasp_rotmat[:, 1]
        grasp_success = self.arm1.move_to_target_pose((target_grasp_pos, target_grasp_quat), duration_s=0.5, settle_s=0.0)

        target_grasp_pos -= 2.4 * target_grasp_rotmat[:, 0]
        grasp_success = self.arm1.move_to_target_pose((target_grasp_pos, target_grasp_quat))
        target_grasp_pos += 0.25 * target_grasp_rotmat[:, 1]
        grasp_success = self.arm1.move_to_target_pose((target_grasp_pos, target_grasp_quat), duration_s=0.5, settle_s=0.0)
        satellite.openPanel1(max_vel=2.0)
        target_grasp_pos -= 0.25 * target_grasp_rotmat[:, 1]
        grasp_success = self.arm1.move_to_target_pose((target_grasp_pos, target_grasp_quat), duration_s=0.5, settle_s=0.0)

        target_grasp_pos += 1.8 * target_grasp_rotmat[:, 2]
        grasp_success = self.arm1.move_to_target_pose((target_grasp_pos, target_grasp_quat))
        target_grasp_pos += 1.2 * target_grasp_rotmat[:, 0]
        target_grasp_pos += 0.5 * target_grasp_rotmat[:, 2]
        grasp_success = self.arm1.move_to_target_pose((target_grasp_pos, target_grasp_quat))
        target_grasp_pos += 0.5 * target_grasp_rotmat[:, 1]
        grasp_success = self.arm1.move_to_target_pose((target_grasp_pos, target_grasp_quat), duration_s=0.5, settle_s=0.0)
        satellite.openAntenna2(max_vel=2.0)
        target_grasp_pos -= 0.5 * target_grasp_rotmat[:, 1]
        grasp_success = self.arm1.move_to_target_pose((target_grasp_pos, target_grasp_quat), duration_s=0.5, settle_s=0.0)

        self.env_runner.sim.setObjectParent(satellite.base_handle, -1, True)
        self.env_runner.sim.resetDynamicObject(satellite.base_handle) 
        self.arm1.move_to_target_pose(self.arm1.home_pose)
        self.arm2.move_to_target_pose(self.arm2.home_pose)
        print("      ✅ 维修完成，已释放目标卫星")

        robot_pos, robot_lin, robot_quat, robot_ang = self.get_state()
        robot_rotmat = R.from_quat(robot_quat).as_matrix()
        target_pos = robot_pos - 5.0 * robot_rotmat[:, 1]
        target_quat = robot_quat
        n_steps = 500
        for step in range(1, n_steps + 1):
            alpha = step / n_steps
            curr_ref_pos = robot_pos + (target_pos - robot_pos) * alpha
            self.step_trajectory_control(curr_ref_pos, target_quat)
            self.env_runner.stepSimulation()
        print("      ✅ 已脱离目标卫星")

    def move_at_satellite_vel(self, duration):
        print(f"\n🟢 [阻塞模式] 跟随目标卫星速度 持续 {duration}s...")
        start_time = self.env_runner.sim.getSimulationTime()
        integral_error = np.zeros(6)
        while self.env_runner.sim.getSimulationTime() - start_time < duration:
            target_vel, target_ang_vel = self.env_runner.sim.getVelocity(self.target_handle)
            target_vel, target_ang_vel = np.array(target_vel), np.array(target_ang_vel)
            target_vel += np.array([target_ang_vel[2], 0.0, -target_ang_vel[0]]) * 7.0
            target_vel_6d = np.concatenate((target_vel, target_ang_vel), axis=0)
            _, _, integral_error = self.step_velocity_control(target_vel_6d, integral_error)
            self.env_runner.stepSimulation()
        print("      ✅ 速度阶段完成")

    def move_with_satellite(self, satellite: SpaceObjectAbstract, duration):
        """
        TODO: Not worked yet
        """

        print(f"\n🟢 [阻塞模式] 跟随目标卫星 持续 {duration}s...")
        start_time = self.env_runner.sim.getSimulationTime()
        while self.env_runner.sim.getSimulationTime() - start_time < duration:
            sat_pos, sat_vel, sat_quat, sat_ang_vel = satellite.get_state()
            sat_rotmat = R.from_quat(sat_quat).as_matrix()
            target_pos = sat_pos - 10.0 * sat_rotmat[:, 1]
            self.step_trajectory_control(target_pos, sat_quat)
            self.env_runner.stepSimulation()
        print("      ✅ 跟随阶段完成")
    
    def plot_results(self, save_dir="figs"):
        """
        生成报告所需的 5 张标准图表
        """
        print("\n📊 正在生成报告专用图表...")
        if not self.data_log: return
        
        if not os.path.exists(save_dir):
            os.makedirs(save_dir)

        # 提取数据
        times = np.array([d['t'] for d in self.data_log])
        phases = np.array([d['phase'] for d in self.data_log])
        
        # 自身状态
        pos = np.array([d['pos'] for d in self.data_log])
        vel = np.array([d['vel'] for d in self.data_log]) # 向量
        ang_vel = np.array([d['ang_vel'] for d in self.data_log]) # 向量
        quat = np.array([d['quat'] for d in self.data_log])
        forces = np.array([d['force'] for d in self.data_log])
        torques = np.array([d['torque'] for d in self.data_log])

        # 目标状态
        t_pos = np.array([d.get('target_pos', np.zeros(3)) for d in self.data_log])
        t_vel = np.array([d.get('target_vel', np.zeros(3)) for d in self.data_log])
        t_ang_vel = np.array([d.get('target_ang_vel', np.zeros(3)) for d in self.data_log])
        t_quat = np.array([d.get('target_quat', [0,0,0,1]) for d in self.data_log])

        # 计算模长
        vel_norm = np.linalg.norm(vel, axis=1)
        t_vel_norm = np.linalg.norm(t_vel, axis=1)
        ang_vel_norm_deg = np.degrees(np.linalg.norm(ang_vel, axis=1))
        t_ang_vel_norm_deg = np.degrees(np.linalg.norm(t_ang_vel, axis=1))

        # 计算相对误差
        # 位置误差
        pos_err = np.linalg.norm(pos - t_pos, axis=1)
        # 姿态误差 (2 * arccos(|q1.dot(q2)|))
        rot_err = []
        for q1, q2 in zip(quat, t_quat):
            dot = np.abs(np.dot(q1, q2))
            dot = np.clip(dot, -1.0, 1.0)
            rot_err.append(2 * np.arccos(dot))
        rot_err = np.array(rot_err)

        # 筛选阶段数据
        # 假设阶段名称为 'approach' 和 'detumble'，根据 main 函数实际设置来
        mask_approach = phases == 'approach'
        mask_detumble = phases == 'detumble'
        
        # 如果没有设置阶段，就用全部数据（fallback）
        if not np.any(mask_approach): mask_approach = np.ones_like(phases, dtype=bool)
        if not np.any(mask_detumble): mask_detumble = np.ones_like(phases, dtype=bool)

        # --- 图 1: 三维空间螺旋逼近轨迹 (Trajectory 3D) ---
        fig1 = plt.figure(figsize=(10, 8))
        ax1 = fig1.add_subplot(111, projection='3d')
        # 绘制 Chaser 轨迹
        ax1.plot(pos[:,0], pos[:,1], pos[:,2], label='Chaser Trajectory', linewidth=2, color='b')
        # 绘制 Target 轨迹
        ax1.plot(t_pos[:,0], t_pos[:,1], t_pos[:,2], label='Target Trajectory', linewidth=2, color='r', linestyle='--')
        # 标记起点和终点
        ax1.scatter(pos[0,0], pos[0,1], pos[0,2], c='g', marker='o', label='Start')
        ax1.scatter(pos[-1,0], pos[-1,1], pos[-1,2], c='k', marker='x', label='End')
        ax1.set_xlabel('X (m)')
        ax1.set_ylabel('Y (m)')
        ax1.set_zlabel('Z (m)')
        ax1.set_title('3D Spiral Trajectory (Inertial Frame)')
        ax1.legend()
        fig1.savefig(os.path.join(save_dir, 'trajectory_3d.png'), dpi=300)
        print("✅ 图1 (trajectory_3d) 已保存")

        # --- 图 2: 状态同步与误差收敛 (Sync Performance) ---
        fig2, (ax2a, ax2b) = plt.subplots(2, 1, figsize=(10, 10), sharex=True)
        
        # 2a: 速度跟随
        ax2a.plot(times, vel_norm, label='|v_chaser|', color='b')
        ax2a.plot(times, t_vel_norm, label='|v_target|', color='r', linestyle='--')
        ax2a.set_ylabel('Linear Speed (m/s)')
        ax2a.legend(loc='upper left')
        ax2a.grid(True)
        
        ax2a_twin = ax2a.twinx()
        ax2a_twin.plot(times, ang_vel_norm_deg, label='|w_chaser|', color='g')
        ax2a_twin.plot(times, t_ang_vel_norm_deg, label='|w_target|', color='orange', linestyle='--')
        ax2a_twin.set_ylabel('Angular Speed (deg/s)')
        ax2a_twin.legend(loc='upper right')

        # 2b: 误差收敛
        ax2b.plot(times, pos_err, label='Pos Error (m)', color='purple')
        ax2b.set_ylabel('Position Error (m)')
        ax2b.legend(loc='upper left')
        ax2b.grid(True)

        ax2b_twin = ax2b.twinx()
        ax2b_twin.plot(times, rot_err, label='Rot Error (rad)', color='brown')
        ax2b_twin.set_ylabel('Rotation Error (rad)')
        ax2b_twin.legend(loc='upper right')
        
        ax2b.set_xlabel('Time (s)')
        ax2a.set_title('State Synchronization & Error Convergence')
        fig2.savefig(os.path.join(save_dir, 'sync_performance.png'), dpi=300)
        print("✅ 图2 (sync_performance) 已保存")

        # --- 图 3: 逼近阶段控制输出 (Control Inputs - Approach) ---
        # 使用 mask_approach 筛选数据
        t_app = times[mask_approach]
        f_app = forces[mask_approach]
        tau_app = torques[mask_approach]

        if len(t_app) > 0:
            fig3, (ax3a, ax3b) = plt.subplots(2, 1, figsize=(10, 8), sharex=True)
            ax3a.plot(t_app, f_app[:,0], label='Fx')
            ax3a.plot(t_app, f_app[:,1], label='Fy')
            ax3a.plot(t_app, f_app[:,2], label='Fz')
            ax3a.set_ylabel('Control Force (N)')
            ax3a.set_title('Control Inputs (Approach Phase)')
            ax3a.legend()
            ax3a.grid(True)

            ax3b.plot(t_app, tau_app[:,0], label='Tx')
            ax3b.plot(t_app, tau_app[:,1], label='Ty')
            ax3b.plot(t_app, tau_app[:,2], label='Tz')
            ax3b.set_ylabel('Control Torque (N·m)')
            ax3b.set_xlabel('Time (s)')
            ax3b.legend()
            ax3b.grid(True)
            fig3.savefig(os.path.join(save_dir, 'control_inputs_approach.png'), dpi=300)
            print("✅ 图3 (control_inputs_approach) 已保存")
        else:
            print("⚠️ 未检测到 Approach 阶段数据，图3跳过")

        # --- 图 4: 复合系统消旋过程 (Detumble Profile) ---
        # 使用 mask_detumble 筛选
        t_det = times[mask_detumble]
        sys_ang_vel_deg = ang_vel_norm_deg[mask_detumble] # 抓取后 chaser 的角速度即为系统的角速度
        sys_ang_vel_vec = np.degrees(ang_vel[mask_detumble])

        if len(t_det) > 0:
            fig4 = plt.figure(figsize=(10, 6))
            plt.plot(t_det, sys_ang_vel_vec[:,0], label='Wx', alpha=0.5)
            plt.plot(t_det, sys_ang_vel_vec[:,1], label='Wy', alpha=0.5)
            plt.plot(t_det, sys_ang_vel_vec[:,2], label='Wz', alpha=0.5)
            plt.plot(t_det, sys_ang_vel_deg, label='|W| (Magnitude)', color='k', linewidth=2)
            plt.title('System Angular Velocity during Detumbling')
            plt.xlabel('Time (s)')
            plt.ylabel('Angular Velocity (deg/s)')
            plt.legend()
            plt.grid(True)
            fig4.savefig(os.path.join(save_dir, 'detumble_profile.png'), dpi=300)
            print("✅ 图4 (detumble_profile) 已保存")

            # --- 图 5: 消旋阶段控制输出 (Control Torque - Detumble) ---
            tau_det = torques[mask_detumble]
            
            fig5 = plt.figure(figsize=(10, 6))
            plt.plot(t_det, tau_det[:,0], label='Tx')
            plt.plot(t_det, tau_det[:,1], label='Ty')
            plt.plot(t_det, tau_det[:,2], label='Tz')
            plt.title('Control Torque during Detumbling')
            plt.xlabel('Time (s)')
            plt.ylabel('Torque (N·m)')
            plt.legend()
            plt.grid(True)
            fig5.savefig(os.path.join(save_dir, 'control_inputs_detumble.png'), dpi=300)
            print("✅ 图5 (control_inputs_detumble) 已保存")
        else:
            print("⚠️ 未检测到 Detumble 阶段数据，图4/5跳过")

        plt.show()