import numpy as np



class PDControllerStableMultiDof(object):
		def __init__(self, pb):
			self._pb = pb

		def computePD(self, bodyUniqueId, jointIndices, desiredPositions, desiredVelocities, kps, kds, maxForces, timeStep):
			
			numJoints = len(jointIndices)#self._pb.getNumJoints(bodyUniqueId)
			curPos,curOrn = self._pb.getBasePositionAndOrientation(bodyUniqueId)
			#q1 = [desiredPositions[0],desiredPositions[1],desiredPositions[2],desiredPositions[3],desiredPositions[4],desiredPositions[5],desiredPositions[6]]
			q1 = [curPos[0],curPos[1],curPos[2],curOrn[0],curOrn[1],curOrn[2],curOrn[3]]
			
			#qdot1 = [0,0,0, 0,0,0,0]
			baseLinVel, baseAngVel = self._pb.getBaseVelocity(bodyUniqueId)
			
			qdot1 = [baseLinVel[0],baseLinVel[1],baseLinVel[2],baseAngVel[0],baseAngVel[1],baseAngVel[2],0]
			qError = [0,0,0, 0,0,0,0]
			
			qIndex=7
			qdotIndex=7
			zeroAccelerations=[0,0,0,  0,0,0,0]
			for i in range (numJoints):
				js = self._pb.getJointStateMultiDof(bodyUniqueId, jointIndices[i])
				
				jointPos=js[0]
				jointVel = js[1]
				q1+=jointPos
				
				if len(js[0])==1:
					desiredPos=desiredPositions[qIndex]
					
					qdiff=desiredPos - jointPos[0]
					qError.append(qdiff)
					zeroAccelerations.append(0.)
					qdot1+=jointVel
					qIndex+=1
					qdotIndex+=1
				if len(js[0])==4:
					desiredPos=[desiredPositions[qIndex],desiredPositions[qIndex+1],desiredPositions[qIndex+2],desiredPositions[qIndex+3]]
					axis = self._pb.getAxisDifferenceQuaternion(desiredPos,jointPos)
					jointVelNew = [jointVel[0],jointVel[1],jointVel[2],0]
					qdot1+=jointVelNew
					qError.append(axis[0])
					qError.append(axis[1])
					qError.append(axis[2])
					qError.append(0)
					desiredVel=[desiredVelocities[qdotIndex],desiredVelocities[qdotIndex+1],desiredVelocities[qdotIndex+2]]
					zeroAccelerations+=[0.,0.,0.,0.]
					qIndex+=4
					qdotIndex+=4
			
			q = np.array(q1)
			qdot=np.array(qdot1)
			
			qdotdesired = np.array(desiredVelocities)
			qdoterr = qdotdesired-qdot
			
			
			Kp = np.diagflat(kps)
			Kd = np.diagflat(kds)
			
			p =  Kp.dot(qError)
			
			#np.savetxt("pb_qError.csv", qError, delimiter=",")
			#np.savetxt("pb_p.csv", p, delimiter=",")
			
			d = Kd.dot(qdoterr)

			#np.savetxt("pb_d.csv", d, delimiter=",")
			#np.savetxt("pbqdoterr.csv", qdoterr, delimiter=",")
			
			
			M1 = self._pb.calculateMassMatrix(bodyUniqueId,q1, flags=1)
			
			
			M2 = np.array(M1)
			#np.savetxt("M2.csv", M2, delimiter=",")
			
			M = (M2 + Kd * timeStep)
			
			#np.savetxt("pbM_tKd.csv",M, delimiter=",")
			
		
			
			c1 = self._pb.calculateInverseDynamics(bodyUniqueId, q1, qdot1, zeroAccelerations, flags=1)
			
			
			c = np.array(c1)
			#np.savetxt("pbC.csv",c, delimiter=",")
			A = M
			#p = [0]*43
			b =  p + d -c
			#np.savetxt("pb_acc.csv",b, delimiter=",")
			qddot = np.linalg.solve(A, b)
			tau = p + d - Kd.dot(qddot) * timeStep
			#print("len(tau)=",len(tau))
			maxF = np.array(maxForces)
			forces = np.clip(tau, -maxF , maxF )
			return forces
			
		
			
class PDControllerStable(object):
		def __init__(self, pb):
			self._pb = pb

		def computePD(self, bodyUniqueId, jointIndices, desiredPositions, desiredVelocities, kps, kds, maxForces, timeStep):
			numJoints = self._pb.getNumJoints(bodyUniqueId)
			jointStates = self._pb.getJointStates(bodyUniqueId, jointIndices)
			q1 = []
			qdot1 = []
			zeroAccelerations = []
			for i in range (numJoints):
				q1.append(jointStates[i][0])
				qdot1.append(jointStates[i][1])
				zeroAccelerations.append(0)
			q = np.array(q1)
			qdot=np.array(qdot1)
			qdes = np.array(desiredPositions)
			qdotdes = np.array(desiredVelocities)
			qError = qdes - q
			qdotError = qdotdes - qdot
			Kp = np.diagflat(kps)
			Kd = np.diagflat(kds)
			p =  Kp.dot(qError)
			d = Kd.dot(qdotError)
			forces = p + d
			
			M1 = self._pb.calculateMassMatrix(bodyUniqueId,q1)
			M2 = np.array(M1)
			M = (M2 + Kd * timeStep)
			c1 = self._pb.calculateInverseDynamics(bodyUniqueId, q1, qdot1, zeroAccelerations)
			c = np.array(c1)
			A = M
			b = -c + p + d
			qddot = np.linalg.solve(A, b)
			tau = p + d - Kd.dot(qddot) * timeStep
			maxF = np.array(maxForces)
			forces = np.clip(tau, -maxF , maxF )
			#print("c=",c)
			return tau