import os
import numpy as np
import pinocchio as pin
from urdf_parser_py.urdf import URDF
import numpy as np


class PIN_utils:
    
    def __init__(self,
                 urdf_path=""):
        # mesh_dir
        mesh_dir = os.path.dirname(os.path.dirname(urdf_path))
        # Load the urdf model
        self.model, self.collision_model, self.visual_model = pin.buildModelsFromUrdf(filename=urdf_path, package_dirs=mesh_dir)
        # Create data required by the algorithms
        self.data = self.model.createData()
  
        print("PIN init !")

    def getModelNames(self,index):
        return self.model.names[index]


    def forwardKinematics(self, q):
        pin.forwardKinematics(self.model, self.data, q)

    def getModelJoints(self,index):
        return self.model.joints[index]

    def getJac(self, q):
        pin.computeJointJacobians(self.model, self.data, q)
        return pin.getJointJacobian(self.model,self.data,7,pin.WORLD)
    
    def getFrameJac(self,q,frameid):
        return pin.computeFrameJacobian(self.model, self.data, q, frameid, pin.WORLD)

    def Jlog6_(self,a):
        return pin.Jlog6(a)
    
    def rpyJac(self,rpy):
        return pin.rpy.computeRpyJacobian(rpy,pin.WORLD)


    def setNumpyArray(self, q):
        return np.asarray(q, dtype=np.float32)
    
    def getInertiaMat(self, q):
        _q = self.setNumpyArray(q)
        pin.crba(self.model, self.data, _q)
        return self.data.M

    
    def forwardABADynamics(self, q, v, tau, flex=None):
        """
        Compute the forward dynamics with contact constraints
        :param  q:  The joint configuration
        :param  v:  The joint velocity
        :param  tau:The joint torque vector
        :return a:  The joint acceleration
        """
        if flex is None:
            return pin.aba(self.model, self.data, q, v, tau)
        else:
            return pin.aba(self.model, self.data, q, v, tau, flex)
    
    def inverseDynamics(self, q, v, a, flex=None):
        """
        The Recursive Newton-Euler algorithm. 
        It computes the inverse dynamics, aka the joint torques according to 
        the current state of the system and the desired joint accelerations.
        :param  q:  The joint configuration
        :param  v:  The joint velocity
        :param  a:  The joint acceleration
        :return tau:The joint torque vector
        """
        if flex is None:
            return pin.rnea(self.model, self.data, q, v, a)
        else:
            return pin.rnea(self.model, self.data, q, v, a, flex)

    def getCoriolisMatrix(self, q, v):
        pin.computeCoriolisMatrix(self.model, self.data, q, v)
        return self.data.C
    
    def getGravity(self, q):
        pin.computeGeneralizedGravity(self.model, self.data, q)
        return self.data.g


    #  computeRNEADerivative
    def getDtau_dq(self, q, v, a):
        pin.computeRNEADerivatives(self.model, self.data, q, v, a)
        return self.data.dtau_dq
    
    def getDtau_dv(self, q, v, a):
        pin.computeRNEADerivatives(self.model, self.data, q, v, a)
        return self.data.dtau_dv

    # def getD2tau_dqdq(self, q, v, a):
    #     # !该函数不存在
    #     d2tau_dqdq = np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0])
    #     d2tau_dvdv = np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0])
    #     dtau_dqdv = np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0])
    #     dtau_dadq = np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0])
    #     pin.ComputeRNEASecondOrderDerivatives(self.model, self.data, q, v, a, d2tau_dqdq, d2tau_dvdv, dtau_dqdv, dtau_dadq)
    #     return self.d2tau_dqdq

    # computeABADerivatives
    def getDdq_dq(self, q, v, tau):
        pin.computeABADerivatives(self.model, self.data, q, v, tau)
        return self.data.ddq_dq
    
    def getDdq_dv(self, q, v, tau):
        pin.computeABADerivatives(self.model, self.data, q, v, tau)
        return self.data.ddq_dv
    
    def getDdq_dtau(self, q, v, tau):
        pin.computeABADerivatives(self.model, self.data, q, v, tau)
        return self.data.ddq_dtau
    
    # computeForwardKinematicsDerivatives
    def getDv_dq(self, q, v, a, joint_id):
        (dv_dq,dv_dv) = pin.getJointVelocityDerivatives(self.model, self.data, joint_id, pin.ReferenceFrame.WORLD) 
        return dv_dq
    
    def getDv_dv(self, q, v, a, joint_id):
        (dv_dq,dv_dv) = pin.getJointVelocityDerivatives(self.model, self.data, joint_id, pin.ReferenceFrame.WORLD) 
        return dv_dv
    
    def getDa_dq(self, q, v, a, joint_id):
        (dv_dq,da_dq,da_dv,da_da) = pin.getJointAccelerationDerivatives(self.model, self.data, joint_id, pin.ReferenceFrame.WORLD)
        return da_dq 
    
    def getDa_dv(self, q, v, a, joint_id):
        (dv_dq,da_dq,da_dv,da_da) = pin.getJointAccelerationDerivatives(self.model, self.data, joint_id, pin.ReferenceFrame.WORLD)
        return da_dv 
    
    def getDa_da(self, q, v, a, joint_id):
        (dv_dq,da_dq,da_dv,da_da) = pin.getJointAccelerationDerivatives(self.model, self.data, joint_id, pin.ReferenceFrame.WORLD)
        return da_da 