import unittest
import pybullet as p
import math, time
import difflib, sys

from utils import allclose, dot


class TestPybulletSaveRestoreMethods(unittest.TestCase):

  def setupWorld(self):
    numObjects = 50

    maximalCoordinates = False

    p.resetSimulation()
    p.setPhysicsEngineParameter(deterministicOverlappingPairs=1)
    p.loadURDF("planeMesh.urdf", useMaximalCoordinates=maximalCoordinates)
    kukaId = p.loadURDF("kuka_iiwa/model_free_base.urdf", [0, 0, 10],
                        useMaximalCoordinates=maximalCoordinates)
    for i in range(p.getNumJoints(kukaId)):
      p.setJointMotorControl2(kukaId, i, p.POSITION_CONTROL, force=0)
    for i in range(numObjects):
      cube = p.loadURDF("cube_small.urdf", [0, i * 0.02, (i + 1) * 0.2])
      #p.changeDynamics(cube,-1,mass=100)
    p.stepSimulation()
    p.setGravity(0, 0, -10)

  def dumpStateToFile(self, file):
    for i in range(p.getNumBodies()):
      pos, orn = p.getBasePositionAndOrientation(i)
      linVel, angVel = p.getBaseVelocity(i)
      txtPos = "pos=" + str(pos) + "\n"
      txtOrn = "orn=" + str(orn) + "\n"
      txtLinVel = "linVel" + str(linVel) + "\n"
      txtAngVel = "angVel" + str(angVel) + "\n"
      file.write(txtPos)
      file.write(txtOrn)
      file.write(txtLinVel)
      file.write(txtAngVel)

  def compareFiles(self, file1, file2):
    diff = difflib.unified_diff(
        file1.readlines(),
        file2.readlines(),
        fromfile='saveFile.txt',
        tofile='restoreFile.txt',
    )
    numDifferences = 0
    for line in diff:
      numDifferences = numDifferences + 1
      sys.stdout.write(line)
    self.assertEqual(numDifferences, 0)

  def testSaveRestoreState(self):
    numSteps = 500
    numSteps2 = 30

    verbose = 0
    self.setupWorld()

    for i in range(numSteps):
      p.stepSimulation()
    p.saveBullet("state.bullet")
    if verbose:
      p.setInternalSimFlags(1)
    p.stepSimulation()
    if verbose:
      p.setInternalSimFlags(0)
      print("contact points=")
      for q in p.getContactPoints():
        print(q)

    for i in range(numSteps2):
      p.stepSimulation()

    file = open("saveFile.txt", "w")
    self.dumpStateToFile(file)
    file.close()

    #################################
    self.setupWorld()

    #both restore from file or from in-memory state should work
    p.restoreState(fileName="state.bullet")
    stateId = p.saveState()

    if verbose:
      p.setInternalSimFlags(1)
    p.stepSimulation()
    if verbose:
      p.setInternalSimFlags(0)
      print("contact points restored=")
      for q in p.getContactPoints():
        print(q)
    for i in range(numSteps2):
      p.stepSimulation()

    file = open("restoreFile.txt", "w")
    self.dumpStateToFile(file)
    file.close()

    p.restoreState(stateId)
    if verbose:
      p.setInternalSimFlags(1)
    p.stepSimulation()
    if verbose:
      p.setInternalSimFlags(0)
      print("contact points restored=")
      for q in p.getContactPoints():
        print(q)
    for i in range(numSteps2):
      p.stepSimulation()

    file = open("restoreFile2.txt", "w")
    self.dumpStateToFile(file)
    file.close()

    file1 = open("saveFile.txt", "r")
    file2 = open("restoreFile.txt", "r")
    self.compareFiles(file1, file2)
    file1.close()
    file2.close()

    file1 = open("saveFile.txt", "r")
    file2 = open("restoreFile2.txt", "r")
    self.compareFiles(file1, file2)
    file1.close()
    file2.close()


#while (p.getConnectionInfo()["isConnected"]):
#	time.sleep(1)

if __name__ == '__main__':
  p.connect(p.DIRECT)
  unittest.main()