#!/usr/bin/env python
from __future__ import print_function
import pyexotica as exo
import math
import time
import numpy as np

PKG = 'exotica_examples'
import roslib; roslib.load_manifest(PKG)  # This line is not needed with Catkin.

import unittest

def get_problem_initializer(collision_scene, URDF):
    return ('exotica/UnconstrainedEndPoseProblem',
            {'Name': 'TestProblem',
             'PlanningScene': [('exotica/Scene',
                                {'CollisionScene': [(collision_scene, {'Name': 'MyCollisionScene'})],
                                 'JointGroup': 'group1',
                                 'Name': 'TestScene',
                                 'Debug': '0',
                                 'SRDF': '{exotica_examples}/test/resources/a_vs_b.srdf',
                                 'SetRobotDescriptionRosParams': '1',
                                 'URDF': URDF})]})

class TestClass(unittest.TestCase):
    def test_continuous_collision(self):
        collision_scene = "exotica/CollisionSceneFCLLatest"

        urdfs_to_test = ['{exotica_examples}/test/resources/primitive_sphere_vs_primitive_sphere_distance.urdf', '{exotica_examples}/test/resources/mesh_vs_mesh_distance.urdf']

        for urdf in urdfs_to_test:
            print("Testing", urdf)
            
            initializer = get_problem_initializer(collision_scene, urdf)
            prob = exo.Setup.create_problem(initializer)
            prob.update(np.zeros(prob.N,))
            scene = prob.get_scene()
            cs = scene.get_collision_scene()

            # Should collide at -2
            p = cs.continuous_collision_check(
                    "A_collision_0", exo.KDLFrame([-3., 0.0, 0.0]), exo.KDLFrame([-1.0, 0.0, 0.0]),
                    "B_collision_0", exo.KDLFrame([0, 0, 0]), exo.KDLFrame([0, 0, 0]))
            np.testing.assert_equal(p.in_collision, True)
            np.testing.assert_array_less((p.time_of_contact - 0.5), 0.1)
            print(p, p.contact_transform_1, p.contact_transform_2)
            np.testing.assert_allclose(p.contact_transform_1.get_translation(), np.array([-2, 0, 0]), atol=0.15)
            np.testing.assert_allclose(p.contact_transform_2.get_translation(), np.array([0, 0, 0]))
            print(p)

if __name__ == '__main__':
    import rostest
    rostest.rosrun(PKG, 'TestContinuousCollision', TestClass)
