def quat_from_axis_angle(x_angle, y_angle, z_angle):
    return quatFromAxisAngle(np.array((x_angle, y_angle, z_angle))) def trans_from_pose(pose):
    return matrixFromPose(pose) def manip_trans_from_object_trans(object_trans, grasp):

    return np.linalg.solve(grasp.T, object_trans.T).T def length(vector):
    return np.linalg.norm(vector) def normalize(vector):
    return 1. / length(vector) * vector def trans_from_point(x, y, z):
    return trans_from_quat_point(unit_quat(), np.array([x, y, z])) def get_point(body):
    return point_from_trans(get_trans(body)) def trans_from_quat(quat):
    return matrixFromQuat(quat) def trans_from_axis_angle(x_angle, y_angle, z_angle):
    return trans_from_quat(quat_from_axis_angle(x_angle, y_angle, z_angle)) def get_active_config(body):
    return body.GetActiveDOFValues() def set_active_config(body, config):
    body.SetActiveDOFValues(config) def set_quat(body, quat):
    set_pose(body, pose_from_quat_point(quat, get_point(body))) def pose_from_quat_point(quat, point):
    return np.concatenate([quat, point]) def irange(start, end, step=1):
    n = start
    while n < end:
        yield n
        n += step INF = float('inf')
from itertools import count
from random import uniform
from openravepy import rotationMatrixFromAxisAngle, Sensor, RaveCreateCollisionChecker, databases, interfaces,  IkParameterization, GeometryType, RaveCreateKinBody, planning_error, openravepy_int,  CollisionOptionsStateSaver, DOFAffine, RaveGetEnvironment, RaveGetAffineDOFValuesFromTransform
from openravepy.misc import SetViewerUserThread
from robotics.openrave.transforms import quat_from_axis_angle, trans_from_pose, manip_trans_from_object_trans,  length, normalize, trans_from_point, get_point,  trans_from_quat, trans_from_axis_angle, get_active_config, set_active_config, set_quat,  pose_from_quat_point
from stripstream.utils import irange, INF
    from motion_planners.rrt_connect import birrt
    from motion_planners.rrt_connect import direct_path
import math
import sys
import numpy as np
        import motion_planners
def set_manipulator_conf(manipulator, values):
    manipulator.GetRobot().SetDOFValues(values, manipulator.GetArmIndices())


def set_base_conf(body, base_values):
    trans = body.GetTransform()
    trans[:3, :3] = rotationMatrixFromAxisAngle(
        np.array([0, 0, base_values[-1]]))
    trans[:2, 3] = base_values[:2]
    body.SetTransform(trans)


def set_gripper(manipulator, values):
    manipulator.GetRobot().SetDOFValues(values, manipulator.GetGripperIndices())


def open_gripper(manipulator):
    _, upper = manipulator.GetRobot().GetDOFLimits(manipulator.GetGripperIndices())
    set_gripper(manipulator, upper)


def close_gripper(manipulator):
    lower, _ = manipulator.GetRobot().GetDOFLimits(manipulator.GetGripperIndices())
    set_gripper(manipulator, lower)


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 openrave_inverse_reachability(ir_model, manip_trans):
    index_manip_iterator = [(manip_trans, manip_trans)]
    try:
        for base_trans, _, _ in ir_model.sampleBaseDistributionIterator(index_manip_iterator, logllthresh=-INF, Nprematuresamples=1):

            yield base_trans
    except (ValueError, planning_error):
        raise StopIteration


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 get_collision_fn(env, body, self_collisions=False):
    def fn(q):
        body.SetActiveDOFValues(q)
        return env.CheckCollision(body) or (self_collisions and body.CheckSelfCollision())
    return fn


def get_sample_fn(env, body, collisions=False, self_collisions=False):
    limits = body.GetActiveDOFLimits()
    collision = get_collision_fn(env, body, self_collisions=self_collisions)

    def fn():
        while True:
            config = np.random.uniform(*limits)
            if not collisions or not collision(config):
                return config
    return fn


def get_distance_fn(body):
    weights = body.GetActiveDOFWeights()

    def distance_fn(q1, q2):
        diff = body.SubtractActiveDOFValues(q2, q1)
        return math.sqrt(np.dot(weights, diff * diff))
    return distance_fn


def get_extend_fn(body):
    resolutions = body.GetActiveDOFResolutions()

    def extend_fn(q1, q2):
        n = int(np.max(
            np.abs(np.divide(body.SubtractActiveDOFValues(q2, q1), resolutions)))) + 1
        q = q1
        for i in range(n):
            q = (1. / (n - i)) * body.SubtractActiveDOFValues(q2, q) + q
            yield q
    return extend_fn


def manip_from_pose_grasp(pose, grasp):
    return manip_trans_from_object_trans(trans_from_pose(pose), grasp)


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)