import pybullet as pb
import time
from pybullet_utils import bullet_client


server = bullet_client.BulletClient(connection_mode=pb.SHARED_MEMORY_SERVER)


print ("Connecting to bullet server")
CONNECTION_METHOD = pb.SHARED_MEMORY
client = bullet_client.BulletClient(connection_mode=CONNECTION_METHOD)


PLANE_PATH = "plane.urdf"
client.loadURDF(PLANE_PATH)

client.setGravity(0, 0, -10)

print ("Adding plane object")
plane_id = client.loadURDF(PLANE_PATH)
print ("Plane ID: %s" % plane_id)

print ("Adding user data to plane")
MyKey1 = client.addUserData(plane_id, "MyKey1", "MyValue1")
MyKey2 = client.addUserData(plane_id, "MyKey2", "MyValue2")
MyKey3 = client.addUserData(plane_id, "MyKey3", "MyValue3")
MyKey4 = client.addUserData(plane_id, "MyKey4", "MyValue4")

print ("Retrieving cached user data")
print (client.getUserData(MyKey1))
print (client.getUserData(MyKey2))
print (client.getUserData(MyKey3))
print (client.getUserData(MyKey4))

print ("Disconnecting")
del client

print ("Reconnecting")
client = bullet_client.BulletClient(connection_mode=CONNECTION_METHOD)

print ("Retrieving synced user data")
print (client.getUserData(MyKey1))
print (client.getUserData(MyKey2))
print (client.getUserData(MyKey3))
print (client.getUserData(MyKey4))

print ("Number of user data entries: %s" % client.getNumUserData(plane_id))

print ("Overriding user data")
client.addUserData(plane_id, "MyKey1", "MyNewValue")

print ("Cached overridden data")
print (client.getUserData(MyKey1))


print ("Disconnecting")
del client

print ("Reconnecting")
client = bullet_client.BulletClient(connection_mode=CONNECTION_METHOD)


print ("Synced overridden data")
print (client.getUserData(MyKey1))

print ("Getting user data ID")
print ("Retrieved ID: %s, ID retrieved from addUserData: %s"  % (client.getUserDataId(plane_id, "MyKey2"), MyKey2))

print ("Removing user data")
client.removeUserData(MyKey2)

print ("Retrieving cached removed data")
print (client.getUserData(MyKey2))

print ("Syncing")
client.syncUserData()

print ("Retrieving removed removed data")
print (client.getUserData(MyKey2))

print ("Iterating over all user data entries and printing results")
for i in range(client.getNumUserData(plane_id)):
  userDataId, key, bodyId, linkIndex, visualShapeIndex = client.getUserDataInfo(plane_id, i)
  print ("Info: (%s, %s, %s, %s, %s)" % (userDataId, key, bodyId, linkIndex, visualShapeIndex))
  print ("Value: %s" % client.getUserData(userDataId))

print ("Removing body")
client.removeBody(plane_id)

print ("Retrieving user data")
print (client.getUserData(MyKey1))
print (client.getUserData(MyKey3))
print (client.getUserData(MyKey4))

print ("Syncing")
client.syncUserData()

print ("Retrieving user data")
print (client.getUserData(MyKey1))
print (client.getUserData(MyKey3))
print (client.getUserData(MyKey4))

plane_id2 = client.loadURDF(PLANE_PATH)
print ("Plane1: %s, plane2: %s" % (plane_id, plane_id2))

print ("Retrieving user data")
print (client.getUserData(MyKey1))
print (client.getUserData(MyKey3))
print (client.getUserData(MyKey4))

