import re
import traceback
import os
import numpy as np
import numpy.testing as nptest

from tesseract_robotics.tesseract_scene_graph import SimpleResourceLocator, SimpleResourceLocatorFn
from tesseract_robotics.tesseract_environment import Environment
from tesseract_robotics.tesseract_common import FilesystemPath, Isometry3d, Translation3d, Quaterniond, \
    ManipulatorInfo
from tesseract_robotics.tesseract_command_language import JointWaypoint, CartesianWaypoint, Waypoint, \
    PlanInstructionType_FREESPACE, PlanInstructionType_START, PlanInstruction, Instruction, \
    isMoveInstruction, isStateWaypoint, CompositeInstruction, flatten, isMoveInstruction, isStateWaypoint
from tesseract_robotics.tesseract_motion_planners import PlannerRequest, PlannerResponse, generateSeed
from tesseract_robotics.tesseract_motion_planners_descartes import DescartesDefaultPlanProfileD, \
    DescartesProblemGeneratorFnD, DescartesMotionPlannerD, DefaultDescartesProblemGeneratorD, DescartesMotionPlannerStatusCategory

def _locate_resource(url):
    try:
        url_match = re.match(r"^package:\/\/tesseract_support\/(.*)$",url)
        if (url_match is None):
            return ""    
        if not "TESSERACT_SUPPORT_DIR" in os.environ:
            return ""
        tesseract_support = os.environ["TESSERACT_SUPPORT_DIR"]
        return os.path.join(tesseract_support, os.path.normpath(url_match.group(1)))
    except:
        traceback.print_exc()

def get_environment():
    locate_resource_fn = SimpleResourceLocatorFn(_locate_resource)
    locator = SimpleResourceLocator(locate_resource_fn)
    env = Environment()
    tesseract_support = os.environ["TESSERACT_SUPPORT_DIR"]
    urdf_path = FilesystemPath(os.path.join(tesseract_support, "urdf/lbr_iiwa_14_r820.urdf"))
    srdf_path = FilesystemPath(os.path.join(tesseract_support, "urdf/lbr_iiwa_14_r820.srdf"))
    assert env.init(urdf_path, srdf_path, locator)
    manip_info = ManipulatorInfo()
    manip_info.manipulator = "manipulator"
    
    return env, manip_info

def test_ompl_freespace_joint_cart():

    env, manip = get_environment()

    fwd_kin = env.getManipulatorManager().getFwdKinematicSolver(manip.manipulator)
    inv_kin = env.getManipulatorManager().getInvKinematicSolver(manip.manipulator)
    joint_names = fwd_kin.getJointNames()
    cur_state = env.getCurrentState()

    wp1 = JointWaypoint(joint_names, np.array([0,0,0,-1.57,0,0,0],dtype=np.float64))
    wp2 = CartesianWaypoint(Isometry3d.Identity() * Translation3d(-.2,.4,0.2) * Quaterniond(0,0,1.0,0))

    start_instruction = PlanInstruction(Waypoint(wp1), PlanInstructionType_START, "TEST_PROFILE")
    plan_f1 = PlanInstruction(Waypoint(wp2), PlanInstructionType_FREESPACE, "TEST_PROFILE")

    program = CompositeInstruction("TEST_PROFILE")
    program.setStartInstruction(Instruction(start_instruction))
    program.setManipulatorInfo(manip)
    program.append(Instruction(plan_f1))

    seed = generateSeed(program, cur_state, env, 3.14, 1.0, 3.14, 10)

    plan_profile = DescartesDefaultPlanProfileD()
    
    test_planner = DescartesMotionPlannerD()
    #test_planner.plan_profiles["TEST_PROFILE"] = plan_profile
    problem_generator_fn = DescartesProblemGeneratorFnD(DefaultDescartesProblemGeneratorD)
    test_planner.problem_generator = problem_generator_fn

    request = PlannerRequest()
    request.seed = seed
    request.instructions = program
    request.env = env
    request.env_state = env.getCurrentState()
    
    response = PlannerResponse()

    assert test_planner.solve(request, response)
    assert response.status.value() == DescartesMotionPlannerStatusCategory.SolutionFound

    results = flatten(response.results)

    assert len(results) == 11
    for instr in results:
        assert isMoveInstruction(instr)
        wp1 = instr.as_MoveInstruction().getWaypoint()
        assert isStateWaypoint(wp1)
        wp = wp1.as_StateWaypoint()
        assert len(wp.joint_names) == 7
        assert isinstance(wp.position,np.ndarray)
        assert len(wp.position) == 7
