def solve_inverse_kinematics(manipulator, manip_trans):
    robot = manipulator.GetRobot()
    env = robot.GetEnv()
    with robot:
        robot.SetActiveDOFs(manipulator.GetArmIndices())

        config = manipulator.FindIKSolution(manip_trans, 0)
        if config is None:
            return None
        robot.SetDOFValues(config, manipulator.GetArmIndices())
        if env.CheckCollision(robot) or robot.CheckSelfCollision():
            return None
        return config def set_manipulator_conf(manipulator, values):
    manipulator.GetRobot().SetDOFValues(values, manipulator.GetArmIndices()) class Conf(Wrapper):
    _ids = count(0) class Traj(Wrapper):
    _ids = count(0) class Grasp(Wrapper):
    _ids = count(0) class Pose(Wrapper):
    _ids = count(0) def manip_from_pose_grasp(pose, grasp):
    return manip_trans_from_object_trans(trans_from_pose(pose), grasp) def manipulator_motion_plan(base_manip, manipulator, goal, step_length=MIN_DELTA, max_iterations=10, max_tries=1):
    with base_manip.robot:
        base_manip.robot.SetActiveManipulator(manipulator)
        base_manip.robot.SetActiveDOFs(manipulator.GetArmIndices())
        with collision_saver(base_manip.robot.GetEnv(), openravepy_int.CollisionOptions.ActiveDOFs):
            try:
                traj = base_manip.MoveManipulator(goal=goal,
                                                  maxiter=max_iterations, execute=False, outputtraj=None, maxtries=max_tries,
                                                  goals=None, steplength=step_length, outputtrajobj=True, jitter=None, releasegil=False)
                return list(sample_manipulator_trajectory(manipulator, traj))
            except planning_error:
                return None def base_motion_plan(base_manip, goal, step_length=MIN_DELTA, max_iterations=10, max_tries=1):

    with base_manip.robot:
        set_active(base_manip.robot, use_base=True)
        with collision_saver(base_manip.robot.GetEnv(), openravepy_int.CollisionOptions.ActiveDOFs):
            try:
                traj = base_manip.MoveActiveJoints(goal=goal, steplength=step_length, maxiter=max_iterations, maxtries=max_tries,
                                                   execute=False, outputtraj=None, goals=None, outputtrajobj=True, jitter=None, releasegil=False, postprocessingplanner=None,
                                                   postprocessingparameters=None)

                return list(sample_base_trajectory(base_manip.robot, traj))
            except planning_error:
                return None def linear_motion_plan(robot, end_config):
    env = robot.GetEnv()
    with robot:

        with collision_saver(env, openravepy_int.CollisionOptions.ActiveDOFs):
            start_config = get_active_config(robot)
            path = [start_config] + \
                list(linear_interpolation(robot, start_config, end_config))
            for conf in path:
                set_active_config(robot, conf)
                if env.CheckCollision(robot):
                    return None
            return path def top_grasps(box):
    (w, l, h) = get_box_dimensions(box)
    origin = trans_from_point(0, 0, -h)
    bottom = trans_from_point(0, 0, -h)
    reflect = trans_from_quat(quat_from_axis_angle(0, -math.pi, 0))
    for i in range(4):
        rotate_z = trans_from_axis_angle(0, 0, i * math.pi / 2)
        yield reflect.dot(origin).dot(bottom).dot(rotate_z) def side_grasps(box, under=True):
    (w, l, h) = get_box_dimensions(box)
    origin = trans_from_point(0, 0, -2 * h)
    for j in range(1 + under):
        swap_xz = trans_from_axis_angle(0, -math.pi / 2 + j * math.pi, 0)
        for i in range(4):
            rotate_z = trans_from_axis_angle(0, 0, i * math.pi / 2)
            yield swap_xz.dot(rotate_z).dot(origin) def random_inverse_reachability(ir_model, manip_trans):
    index_manip_iterator = [(manip_trans, manip_trans)]
    try:
        for base_trans, _, _ in ir_model.randomBaseDistributionIterator(index_manip_iterator):

            yield base_trans
    except (ValueError, planning_error):
        raise StopIteration def mp_birrt(robot, q1, q2):
    from motion_planners.rrt_connect import birrt
    return birrt(q1, q2, get_distance_fn(robot),
                 get_sample_fn(robot.GetEnv(), robot),
                 get_extend_fn(robot), get_collision_fn(robot.GetEnv(), robot)) def set_active(robot, indices=(), use_base=False):
    if use_base is False:
        robot.SetActiveDOFs(indices)
    robot.SetActiveDOFs(indices, AFFINE_MASK, ROTATION_AXIS) def set_pose(body, pose):
    set_trans(body, trans_from_pose(pose)) def object_trans_from_manip_trans(manip_trans, grasp):
    return np.dot(manip_trans, grasp) def trans_from_point(x, y, z):
    return trans_from_quat_point(unit_quat(), np.array([x, y, z])) def set_trans(body, trans):
    body.SetTransform(trans) def set_full_config(body, config, dof_indices=None):
    if dof_indices is None:
        body.SetConfigurationValues(config)
    else:
        full_config = get_full_config(body)
        full_config[dof_indices] = config
        body.SetConfigurationValues(full_config) def base_values_from_full_config(config):
    return base_values_from_pose(pose_from_full_config(config)) def full_config_from_base_values(base_values, config):
    _, _, z = point_from_full_config(config)
    return full_config_from_pose(pose_from_base_values(base_values, z=z), config)
from itertools import islice
from robotics.openrave.utils import solve_inverse_kinematics,  set_manipulator_conf, Conf, Traj, Grasp, Pose, manip_from_pose_grasp,  manipulator_motion_plan, base_motion_plan, linear_motion_plan, top_grasps, side_grasps, random_inverse_reachability,  mp_birrt, set_active
from robotics.openrave.transforms import set_pose,  object_trans_from_manip_trans, trans_from_point, set_trans, set_full_config, base_values_from_full_config,  full_config_from_base_values
from time import sleep
import numpy as np


def get_stable_test(bodies, surfaces):
    def stable_test(obj, pose, surface):
        bodies[obj].Enable(True)
        set_pose(bodies[obj], pose.value)
        return surfaces[surface].supports(bodies[obj])
    return stable_test


def base_generator_fn(ir_model, max_failures=20):
    env = ir_model.env
    robot = ir_model.robot

    def base_generator(obj, pose, grasp):

        manip_tform = manip_from_pose_grasp(pose.value, grasp.value)
        generator = random_inverse_reachability(ir_model, manip_tform)
        while True:
            for _ in range(max_failures):
                base_tform = next(generator)
                if base_tform is None:
                    continue
                set_trans(robot, base_tform)
                if env.CheckCollision(robot) or robot.CheckSelfCollision():
                    continue
                yield base_tform
                break
            else:
                return
    return base_generator


def grasp_generator_fn(bodies, use_top_grasps, use_side_grasps, max_grasps):
    def grasp_generator(obj):
        grasps = []
        if use_top_grasps:
            grasps += list(top_grasps(bodies[obj]))
        if use_side_grasps:
            grasps += list(side_grasps(bodies[obj]))
        yield [(Grasp(grasp),) for grasp in grasps[:max_grasps]]
    return grasp_generator


def pose_generator_fn(bodies, surfaces):
    def pose_generator(obj, surface):
        while True:
            bodies[obj].Enable(True)
            pose = surfaces[surface].sample_placement(bodies[obj])
            if pose is None:
                break
            yield [(Pose(pose),)]
    return pose_generator


def cfree_pose_fn(env, bodies):

    def cfree_pose(obj1, pose1, obj2, pose2):
        body1 = bodies[obj1]
        body1.Enable(True)
        set_pose(body1, pose1.value)

        body2 = bodies[obj2]
        body2.Enable(True)
        set_pose(body2, pose2.value)
        return not env.CheckCollision(body1, body2)
    return cfree_pose


"""
def cfree_traj_fn(env, robot, manipulator, body1, body2, all_bodies):
  def _cfree_traj_pose(traj, pose): # Collision free test between a robot executing traj and an object at pose
    enable_all(all_bodies, False)
    body2.Enable(True)
    set_pose(body2, pose.value)
    for conf in traj.value:
      set_manipulator_conf(manipulator, conf)
      if env.CheckCollision(robot, body2):
        return False
    return True

  def _cfree_traj_grasp_pose(traj, grasp, pose): # Collision free test between an object held at grasp while executing traj and an object at pose
    enable_all(all_bodies, False)
    body1.Enable(True)
    body2.Enable(True)
    set_pose(body2, pose.value)
    for conf in traj.value:
      set_manipulator_conf(manipulator, conf)
      manip_trans = manipulator.GetTransform()
      set_pose(body1, object_trans_from_manip_trans(manip_trans, grasp.value))
      if env.CheckCollision(body1, body2):
        return False
    return True

  def cfree_traj(traj, pose): # Collision free test between a robot executing traj (which may or may not involve a grasp) and an object at pose
    if DISABLE_MOTION_COLLISIONS:
      return True
    if traj.pose is not None and traj.pose == pose:
      # This is the same pose of the manipulation
      return True
    return _cfree_traj_pose(traj, pose) and (traj.grasp is None or _cfree_traj_grasp_pose(traj, traj.grasp, pose))
  return cfree_traj
"""


def sample_grasp_traj_fn(ir_model, bodies, carry_config, max_failures=100):
    env = ir_model.env
    manipulator = ir_model.manip

    robot = ir_model.robot

    def sample_grasp_traj(obj, pose, grasp):
        enable_all(bodies, False)
        body = bodies[obj]
        body.Enable(True)
        set_pose(body, pose.value)

        set_manipulator_conf(manipulator, carry_config)
        manip_tform = manip_from_pose_grasp(pose.value, grasp.value)
        for base_tform in islice(random_inverse_reachability(ir_model, manip_tform), max_failures):

            set_trans(robot, base_tform)