from ruckig import InputParameter, Ruckig, Trajectory, Result
import pinocchio as pin
import numpy as np
from time import time
import copy

def get_joint_id(model, name):
    for i, n in enumerate(model.names):
        if n == name:
            return i

def get_joint_index_q(model, name):
    for i, n in enumerate(model.names):
        if n == name:
            return model.joints[i].idx_q
        
def select_pos_from_positions(model, positions, target_positions):
    resp = {}
    for joint_name in positions:
        resp[joint_name] = target_positions[get_joint_index_q(model, joint_name)]
    return resp

def quaternion_to_rotation_matrix(q):
    """
    Convert a quaternion to a 3x3 rotation matrix.

    Parameters:
        q (array-like): Quaternion [x, y, z, w]  # Note the order change

    Returns:
        numpy.ndarray: 3x3 rotation matrix
    """
    x, y, z, w = q
    R = np.array([
        [1 - 2*y**2 - 2*z**2, 2*x*y - 2*w*z, 2*x*z + 2*w*y],
        [2*x*y + 2*w*z, 1 - 2*x**2 - 2*z**2, 2*y*z - 2*w*x],
        [2*x*z - 2*w*y, 2*y*z + 2*w*x, 1 - 2*x**2 - 2*y**2]
    ])
    return R

def rotation_matrix_to_quaternion(R):
    """
    Convert a 3x3 rotation matrix to a quaternion.

    Parameters:
        R (numpy.ndarray): 3x3 rotation matrix
    """
    return pin.Quaternion(R).coeffs()

class Kinematic_dynamic_planner():

    def __init__(self, 
                 model: pin.Model,
                 max_velocity: np.ndarray = None, max_acceleration: np.ndarray = None, max_jerk: np.ndarray = None):

        ## Pinocchio mode ##
        self.model = model
        self.data = self.model.createData()

        if max_velocity is None: max_velocity = self.model.velocityLimit
        if max_acceleration is None: max_acceleration = max_velocity*10
        if max_jerk is None: max_jerk = max_acceleration*10
        
        assert len(max_velocity) == len(max_acceleration) == len(max_jerk), "input constraints should be the same size"
        self.nDOF = len(max_velocity)

        ## Ruckig planner ##
        self.constraints = InputParameter(self.nDOF)
        self.otg = Ruckig(self.nDOF)
        self.trajectory = Trajectory(self.nDOF)


        self.init_constraints(max_velocity, max_acceleration, max_jerk)


    def init_constraints(self, max_velocity, max_acceleration, max_jerk):

        # Set current and target states to 0
        self.constraints.current_position = [0.0]*self.nDOF
        self.constraints.current_velocity = [0.0]*self.nDOF
        self.constraints.current_acceleration = [0.0]*self.nDOF

        self.constraints.target_position = [0.0]*self.nDOF
        self.constraints.target_velocity = [0.0]*self.nDOF
        self.constraints.target_acceleration = [0.0]*self.nDOF

        # Set limits
        self.constraints.max_velocity = list(max_velocity)
        self.constraints.max_acceleration = list(max_acceleration)
        self.constraints.max_jerk = list(max_jerk)


    def set_joint_target(self,
                         current_position, current_velocity,
                         target_position, target_velocity):

            '''
            Set the current and target joint position and velocities

            Input: current_position, current_velocity, target_position, target_velocity: np.ndarray or list
            Output: None
            '''
            
            self.t0 = time()
            
            # assert len(current_position) == len(current_velocity) == len(target_position) == len(target_velocity), "input constraints should be the same size"

            self.constraints.current_position = list(current_position)
            self.constraints.current_velocity = list(current_velocity)

            self.constraints.target_position = list(target_position)
            self.constraints.target_velocity = list(target_velocity)

            result = self.otg.calculate(self.constraints, self.trajectory)


    def set_endEffector_target(self, target_joint_id, current_joint_position, current_joint_velocity,
                                    target_position, target_rotation, target_twist):
             
            '''
            Set the end-effector target pose and twist
    
            Input: target_pose, target_twist: np.ndarray or list
            Output: None
            '''
    
            oMdes = pin.SE3(target_rotation, target_position)
  
            q      = current_joint_position.copy()
            eps    = 1e-4
            IT_MAX = 100
            DT     = 1e-1
            JOINT_ID = target_joint_id # self.nDOF
            
            for _ in range(IT_MAX):
                pin.forwardKinematics(self.model, self.data, q)
                dMi = oMdes.actInv(self.data.oMi[JOINT_ID])
                # dMi = self.data.oMi[JOINT_ID].actInv(oMdes)
                err = pin.log(dMi).vector

                if np.linalg.norm(err) < eps:
                    # print('converged', err, np.linalg.norm(err))
                    break

                J = pin.computeJointJacobian(self.model, self.data, q, JOINT_ID)
                v = - J.T.dot(np.linalg.solve(J.dot(J.T) + 1e-12 * np.eye(6), err))
                q = pin.integrate(self.model, q, v*DT)
                q = np.clip(q, self.model.lowerPositionLimit, self.model.upperPositionLimit)

            # Compute final Jacobian after IK convergence
            pin.forwardKinematics(self.model, self.data, q)
            J = pin.computeJointJacobian(self.model, self.data, q, JOINT_ID)
            
            target_qdot = np.linalg.pinv(J).dot(target_twist)

            self.set_joint_target(current_joint_position, current_joint_velocity, q, target_qdot)

    def forward_kinematics(self, target_joint_id, joint_position, joint_velocity, format="quaternion"):
             
            '''
            Compute the end-effector pose and twist from the joint position
    
            Input: joint_position: np.array
            Output: 
            - end_effector_translation: np.array
            - end_effector_rotation: np.array
            - end_effector_twist: np.array
            '''

            assert format in ["quaternion", "rotation_matrix"], "format should be either 'quaternion' or 'rotation_matrix'"

            pin.forwardKinematics(self.model, self.data, joint_position)
            J = pin.computeJointJacobian(self.model, self.data, joint_position, target_joint_id)

            end_effector_translation = copy.copy(self.data.oMi[target_joint_id].translation)
            end_effector_rotation = copy.copy(self.data.oMi[target_joint_id].rotation)
            if format == "quaternion":
                end_effector_rotation = pin.Quaternion(end_effector_rotation).coeffs()

            end_effector_twist = J@joint_velocity
    
            return end_effector_translation, end_effector_rotation, end_effector_twist

    def get_joint_target(self):

        '''
        Compute the target joint position, velocity, and torque at a given time
        The joint state is sampled from the trajectory at the current time
        and the feedforward torque is computed using the inverse dynamics model

        Input: None
        Output: target_position, target_velocity, target_torque: np.array
        '''

        current_time = time() - self.t0

        if current_time > self.trajectory.duration: 
            # print("Trajectory finished, tracking target position")
            current_time = self.trajectory.duration

        target_position, target_velocity, target_acceleration = self.trajectory.at_time(current_time)
        target_position = np.array(target_position)
        target_velocity = np.array(target_velocity)
        target_acceleration = np.array(target_acceleration)

        self.torque = pin.rnea(self.model, self.data, 
                                 target_position, target_velocity, target_acceleration)
        
        return target_position, target_velocity, self.torque

class EasyIKPlanner():

    def __init__(self, model: pin.Model):

        MAX_ACCELERATION = 0.3
        MAX_JERK = 0.05

        self.model = model
        self.q = pin.neutral(self.model)

        self.planner = Kinematic_dynamic_planner(model, max_acceleration=[MAX_ACCELERATION]*23, max_jerk=[MAX_JERK]*23)

    def update_q(self, q: np.ndarray):
        self.q = q

    def ik(self, pose, is_right: bool = True):
        joint_map = ["SpR", "SrR", "SwR", "EpR", "WwR", "WpR", "WrR"]
        target_joint_name = "WrR"
        if not is_right:
            joint_map = ["SpL", "SrL", "SwL", "EpL", "WwL", "WpL", "WrL"]
            target_joint_name = "WrL"
        
        self.planner.set_endEffector_target(get_joint_id(self.model, target_joint_name), current_joint_position=self.q, current_joint_velocity=np.zeros(self.model.nv),
            target_position=pose[:3], target_twist=np.zeros(6), 
            target_rotation=quaternion_to_rotation_matrix(pose[3:]))
        arm = select_pos_from_positions(self.model, joint_map, self.planner.constraints.target_position)
        return arm

