from pinocchio.utils import *
from pinocchio.explog import exp,log
from numpy.linalg import pinv,norm
from pinocchio import forwardKinematics, Inertia, JointModelRX, Model, SE3
import gepetto.corbaserver
from display import Display
import eigenpy
eigenpy.switchToNumpyArray()

class Visual:
    def __init__(self,name,jointParent,placement):
        self.name = name                  # Name in gepetto viewer
        self.jointParent = jointParent    # ID (int) of the joint 
        self.placement = placement        # placement of the body wrt joint, i.e. bodyMjoint
    def place(self,display,oMjoint):
        oMbody = oMjoint*self.placement
        display.place(self.name,oMbody,False)
class Leg:
    def __init__(self):
        self.viewer = Display()
        self.visuals = []
        self.model = Model ()
        self.CreateLeg()
        self.data = self.model.createData()
        self.q0 = zero(self.model.nq)

    def CreateLeg(self,rootId=0, prefix='', jointPlacement=SE3.Identity()):
        color   = [red,green,blue,transparency] = [1,1,0.78,1.0]
        colorred = [1.0,0.0,0.0,1.0]

        jointId = rootId

        jointId = rootId

        jointName          = prefix + "hip_joint"
        joint              = JointModelRX()
        jointId = self.model.addJoint(jointId,joint,jointPlacement,jointName)
        self.model.appendBodyToJoint(jointId,Inertia.Random(),SE3.Identity())
        self.viewer.viewer.gui.addSphere('world/' + prefix + 'hip', 0.3,colorred)
        self.visuals.append( Visual('world/' + prefix + 'hip',jointId,SE3.Identity()) )
        self.viewer.viewer.gui.addBox('world/' + prefix + 'thigh', .1,.1,.5,color)
        self.visuals.append( Visual('world/' + prefix + 'thigh',jointId,SE3(eye(3),np.array([0., 0., -.5]))))

        jointName          = prefix + "knee_joint"
        jointPlacement     = SE3(eye(3),np.array( [0, 0, -1.0] ))
        joint              = JointModelRX()
        jointId = self.model.addJoint(jointId,joint,jointPlacement,jointName)
        self.model.appendBodyToJoint(jointId,Inertia.Random(),SE3.Identity())
        self.viewer.viewer.gui.addSphere('world/' + prefix + 'knee', 0.3,colorred)
        self.visuals.append( Visual('world/' + prefix + 'knee',jointId,SE3.Identity()) )
        self.viewer.viewer.gui.addBox('world/' + prefix + 'tibia', .1,.1,.5,color)
        self.visuals.append( Visual('world/' + prefix + 'tibia',jointId,SE3(eye(3),np.array([0., 0., -.5]))))

        jointName          = prefix + "ankle_joint"
        jointPlacement     = SE3(eye(3),np.array( [0, 0, -1.0] ))
        joint              = JointModelRX()
        jointId = self.model.addJoint(jointId,joint,jointPlacement,jointName)
        self.model.appendBodyToJoint(jointId,Inertia.Random(),SE3.Identity())
        self.viewer.viewer.gui.addSphere('world/' + prefix + 'ankle', 0.3,colorred)
        self.visuals.append( Visual('world/' + prefix + 'ankle',jointId,SE3.Identity()) )
        self.viewer.viewer.gui.addBox('world/' + prefix + 'feet', .25,.5,.1,color)
        self.visuals.append( Visual('world/' + prefix + 'feet',jointId,SE3(eye(3),np.array([0., 0., -.25]))))

    def display(self,q):
        forwardKinematics(self.model,self.data,q)
        for visual in self.visuals:
            visual.place( self.viewer,self.data.oMi[visual.jointParent] )
        self.viewer.viewer.gui.refresh()