from pinocchio.utils import *
from pinocchio.explog import exp,log
from numpy.linalg import pinv,norm
from pinocchio import forwardKinematics, Inertia, JointModelRX, Model, neutral, SE3,JointModelPX,JointModelPZ
import gepetto.corbaserver
from display import Display
import numpy as np
import eigenpy
import time
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 Ball:
    def __init__(self,angle=np.pi/3,speed=5):
        self.angle=angle
        self.speed=speed
        
        self.viewer = Display()
        self.visuals = []
        self.model = Model ()
        self.Create()
        self.data = self.model.createData()
        self.Throw()
        self.q0 = zero(self.model.nq) 

    def Create(self,rootId=0, 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

        jointName          = "BallX"
        jointPlacement     = SE3(eye(3),np.array( [0, 0, 2] ))
        joint              = JointModelPX()
        jointId = self.model.addJoint(jointId,joint,jointPlacement,jointName)
        #self.model.appendBodyToJoint(jointId,Inertia.Random(),SE3.Identity())
        
        jointName          = "BallZ"
        jointPlacement     = SE3(eye(3),np.array( [0, 0, 0] ))
        joint              = JointModelPZ()
        jointId = self.model.addJoint(jointId,joint,jointPlacement,jointName)
        self.model.appendBodyToJoint(jointId,Inertia.Random(),SE3.Identity())
        
        self.viewer.viewer.gui.addSphere('world/' +'ball', 0.3,colorred)
        self.visuals.append( Visual('world/' +'ball',jointId,SE3.Identity()) )
        
    def Throw(self):
        dt = 1e-3
        q = neutral (self.model)
        v = np.array (self.model.nv * [0.0])
        v1 = np.array (self.model.nv * [0.0])
        v1[1]=-0.0098
        v[0]=self.speed*np.cos(self.angle)
        v[1]=self.speed*np.sin(self.angle)
        print(v)
        while (q[1]>-2):
            q+=v*dt
            v+=v1
            self.display(q)
            time.sleep(10*dt)
        
    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()