from pinocchio.utils import *
from pinocchio.explog import exp,log
from numpy.linalg import pinv,norm
from pinocchio import forwardKinematics, Inertia, JointModelRX,JointModelPY,JointModelPZ, Model, SE3,neutral
import gepetto.corbaserver
import time
from display import Display
import eigenpy
import numpy as np
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 Robot():
    def __init__(self):
        self.thigh = 1.
        self.tibia = 1.
        self.feetsize = (.2,.5,.1)
        self.feet = 0.2
        self.leglen=1.+1.+.3
        self.legdist=0.5

        self.viewer = Display()
        self.visuals = []
        self.model = Model ()
        self.Create()
        self.data = self.model.createData()
        self.q0 = zero(self.model.nq)
        #print(self.model)
        #self.Walk(10)

    def CreateLeg(self,rootId=0,name='',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          = name+"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/'  +name+ 'hip', 0.1,colorred)
        self.visuals.append( Visual('world/'  + name+'hip',jointId,SE3.Identity()) )
        self.viewer.viewer.gui.addBox('world/'  +name+ 'thigh', .1,.1,.5,color)
        self.visuals.append( Visual('world/' + name+'thigh',jointId,SE3(eye(3),np.array([0., 0., -.5]))))

        jointName          = name+"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/'  +name+ 'knee', 0.1,colorred)
        self.visuals.append( Visual('world/'  + name+'knee',jointId,SE3.Identity()) )
        self.viewer.viewer.gui.addBox('world/'  +name+ 'tibia', .1,.1,.5,color)
        self.visuals.append( Visual('world/'  +name+ 'tibia',jointId,SE3(eye(3),np.array([0., 0., -.5]))))

        jointName          = name+"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/'  + name+'ankle', 0.1,colorred)
        self.visuals.append( Visual('world/' + name+'ankle',jointId,SE3.Identity()) )
        self.viewer.viewer.gui.addBox('world/' +name+ 'feet', .25,.5,.1,color)
        self.visuals.append( Visual('world/' +name+ 'feet',jointId,SE3(eye(3),np.array([0., .1, -.25]))))

    def Crleft(self):
        self.CreateLeg(name='left_',jointPlacement=SE3(np.eye(3), np.array([self.legdist/2,0, self.leglen])))     
    def Crright(self):
        self.CreateLeg(name='right_',jointPlacement=SE3(np.eye(3), np.array([-self.legdist/2,0, self.leglen])))
    def Create(self):
        self.Crleft()
        self.Crright()
    
                    
    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()