#!/usr/bin/env python

import roslib
import unittest
import pyexotica as exo
PKG = 'exotica_examples'
roslib.load_manifest(PKG)  # This line is not needed with Catkin.


def test_scene(problem_name, clean_scene=False):
    print('Loading '+problem_name)
    (sol, prob) = exo.Initializers.load_xml_full(
        '{exotica_examples}/test/resources/test_scene.xml', problem_name=problem_name)
    problem = exo.Setup.create_problem(prob)
    solver = exo.Setup.create_solver(sol)
    solver.specify_problem(problem)

    print('Solving original scene ...')
    problem.update(problem.start_state)

    print('Updating the scene ...')
    if clean_scene:
        problem.get_scene().clean_scene()
        problem.get_scene().load_scene_file(
            '{exotica_examples}/resources/scenes/example_moving_obstacle.scene')
    problem.get_scene().load_scene_file(
        '{exotica_examples}/resources/scenes/example_distance.scene')
    print('Solving with updated scene ...')

    problem.update(problem.start_state)
    print('Done')


class TestClass(unittest.TestCase):
    def test_MinimalProblem(self):
        test_scene('MinimalProblem')

    def test_TrajectoryProblem(self):
        test_scene('TrajectoryProblem')

    def test_CollisionProblem(self):
        test_scene('CollisionProblem')

    def test_CustomLink1Problem(self):
        test_scene('CustomLink1Problem')

    def test_CustomLink2Problem(self):
        test_scene('CustomLink2Problem')

    def test_CustomLink3Problem(self):
        test_scene('CustomLink3Problem')

    def test_FullProblem(self):
        test_scene('FullProblem')

    def test_clean_MinimalProblem(self):
        test_scene('MinimalProblem', True)

    def test_clean_TrajectoryProblem(self):
        test_scene('TrajectoryProblem', True)

    def test_clean_CollisionProblem(self):
        test_scene('CollisionProblem', True)

    def test_clean_CustomLink1Problem(self):
        test_scene('CustomLink1Problem', True)

    def test_clean_CustomLink2Problem(self):
        test_scene('CustomLink2Problem', True)

    def test_clean_CustomLink3Problem(self):
        test_scene('CustomLink3Problem', True)

    def test_clean_FullProblem(self):
        test_scene('FullProblem', True)

    def test_creation(self):
        scene_conf = {'URDF': "{exotica_examples}/resources/robots/lwr_simplified.urdf",
                      'SRDF': "{exotica_examples}/resources/robots/lwr_simplified.srdf",
                      'JointGroup': "arm"}

        # create scene with URDF frames
        scene1 = exo.Setup.create_scene(('exotica/Scene', scene_conf))
        scene_links1 = scene1.get_tree_names()

        # add additional frame
        frame_name = "camera_frame"
        scene_conf.update({'Links': [('exotica/Link', {'Name': frame_name, 'Transform': [0, 0, 0, 1, 0, 0, 0]})]})
        scene2 = exo.Setup.create_scene(('exotica/Scene', scene_conf))
        scene_links2 = scene2.get_tree_names()

        # check
        self.assertEqual((set(scene_links1) | {frame_name}), set(scene_links2))


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