"""
Robot Data Structure
Python equivalent of the C++ RobotData class
"""
import numpy as np
from scipy.spatial.transform import Rotation
from common.joystick import ControlFlag
import copy

class RobotData:
    """机器人状态数据结构"""

    def __init__(self, motor_num: int = 29, whole_joint_num: int = 35):
        self.motor_num = motor_num
        self.whole_joint_num = whole_joint_num

        # Joint states (actual)
        self.q_a_ = np.zeros(whole_joint_num)  # Joint positions
        self.q_dot_a_ = np.zeros(whole_joint_num)  # Joint velocities
        self.tau_a_ = np.zeros(whole_joint_num)  # Joint torques
        self.temperature_a = np.zeros(motor_num)  # Joint temperatures
        self.q_a_last = np.zeros(whole_joint_num) # 上一时刻关节位置
        self.qdot_a_last = np.zeros(whole_joint_num) # 上一时刻关节速度
        self.tor_a_last = np.zeros(whole_joint_num) # 上一时刻关节力矩

        # Joint commands (desired)
        self.q_d_ = np.zeros(whole_joint_num)  # Desired joint positions
        self.q_d_s_ = np.zeros(whole_joint_num) # Desired serial joint positions 
        self.q_dot_d_ = np.zeros(whole_joint_num)  # Desired joint velocities
        self.tau_d_ = np.zeros(whole_joint_num)  # Desired joint torques

        # Control gains
        self.joint_kp_p_ = np.zeros(motor_num)  # Proportional gains
        self.joint_kd_p_ = np.zeros(motor_num)  # Derivative gains

        # IMU data: [yaw, pitch, roll, gyro_x, gyro_y, gyro_z, acc_x, acc_y, acc_z]
        self.imu_data_ = np.zeros(13)
        self.imu_quat_ = np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float32)  # w, x, y, z

        # Timing
        self.time_now_ = 0.0

        # Configuration
        self.config_file_ = ""

        # control cmd
        # walk command
        self.walk_cmd_ = np.zeros(3) # x_speed, y_speed, yaw_speed
        # 控制符
        self.control_flag = ControlFlag()


    def copy_from(self, other: 'RobotData'):
        """从另一个RobotData对象复制数据"""
        self.q_a_[:] = other.q_a_[:]
        self.q_dot_a_[:] = other.q_dot_a_[:]
        self.tau_a_[:] = other.tau_a_[:]
        self.q_d_[:] = other.q_d_[:]
        self.q_dot_d_[:] = other.q_dot_d_[:]
        self.tau_d_[:] = other.tau_d_[:]
        self.joint_kp_p_[:] = other.joint_kp_p_[:]
        self.joint_kd_p_[:] = other.joint_kd_p_[:]
        self.imu_data_[:] = other.imu_data_[:]
        self.imu_quat_[:] = other.imu_quat_[:]
        self.time_now_ = other.time_now_
        self.config_file_ = other.config_file_
        self.walk_cmd_[:] = other.walk_cmd_[:]
        self.control_flag = copy.deepcopy(other.control_flag)


    def get_joint_pos(self) -> np.ndarray:
        joint_start_idx = 35 - self.motor_num
        joint_pos = self.q_a_[joint_start_idx:].astype(np.float32) 
        return joint_pos
    
    def get_serial_joint_pos_desired(self) -> np.ndarray:
        joint_start_idx = 35 - self.motor_num
        joint_pos_desired = self.q_d_s_[joint_start_idx:].astype(np.float32)
        return joint_pos_desired

    def get_joint_vel(self)-> np.ndarray:
        joint_start_idx = 35 - self.motor_num
        joint_vel = self.q_dot_a_[joint_start_idx:].astype(np.float32)
        return joint_vel

    def get_angular_velocity(self) -> np.ndarray:
        omega_xyz = np.array([
            self.imu_data_[3],
            self.imu_data_[4],
            self.imu_data_[5]
        ], dtype=np.float32)
        return omega_xyz
    
    def get_robot_quat(self):
        if np.linalg.norm(self.imu_quat_) > 0.0:
            return self.imu_quat_.astype(np.float32)
        rpy = np.array([
            self.imu_data_[2],  # roll
            self.imu_data_[1],  # pitch
            self.imu_data_[0]   # yaw
        ], dtype=np.float32) * 1.0    
        robot_quat_wxyz = self.euler_to_quaternion_scipy(rpy[0], rpy[1], rpy[2])
        return robot_quat_wxyz

    def euler_to_quaternion_scipy(self, roll, pitch, yaw, degrees=False):
        """
        使用SciPy进行欧拉角转四元数
        参数:
            roll: 绕x轴的旋转角度
            pitch: 绕y轴的旋转角度  
            yaw: 绕z轴的旋转角度
            degrees: 输入角度是否为度，默认为弧度
        返回:
            [w, x, y, z]: 四元数分量 (w为实部)
        """
        # 创建旋转对象 (顺序: 'xyz' 对应 roll, pitch, yaw)
        rotation = Rotation.from_euler('xyz', [roll, pitch, yaw], degrees=degrees)
        
        # 转换为四元数 (顺序: [x, y, z, w])
        quaternion = rotation.as_quat()
    
        return [quaternion[3], quaternion[0], quaternion[1], quaternion[2]]  # 返回 w, x, y, z
    
    def get_waist_yrp(self) -> np.ndarray:
        joint_pos  = self.get_joint_pos()
        waist_yaw, waist_roll, waist_pitch = joint_pos[12], joint_pos[13], joint_pos[14]
        return np.array([waist_yaw, waist_roll, waist_pitch], dtype=np.float32)
    
    def get_base_linear_acceleration(self) -> np.ndarray:
        lin_acc = np.array([
            self.imu_data_[6],
            self.imu_data_[7],
            self.imu_data_[8]
        ], dtype=np.float32)
        return lin_acc
    
    def get_project_gravity(self) -> np.ndarray:
        """根据机器人姿态重力投影(待完善)

        Args:
            None
        """        
        robot_quat_wxyz = self.get_robot_quat()
        robot_quat_xyzw = np.array([robot_quat_wxyz[1], robot_quat_wxyz[2], robot_quat_wxyz[3], robot_quat_wxyz[0]])
        g = np.array([0., 0., -1.])
        projected_gravity = self.quat_rotate_inverse_numpy(robot_quat_xyzw, g)
        return projected_gravity
    
    def quat_rotate_inverse_numpy(self, q, v):
        """
        q: [x, y, z, w], shape=(4)\\
        v: [x, y, z], shape=(3)
        """
        q_w = q[3]
        q_vec = q[:3]
        a = v * (2.0 * q_w ** 2 - 1.0)
        b = np.cross(q_vec, v) * q_w * 2.0
        c = q_vec * np.dot(q_vec, v) * 2.0
        return a - b + c

    def get_walk_cmd(self) -> np.ndarray:
        """获取行走命令: [x_speed, y_speed, yaw_speed]"""
        return self.walk_cmd_