#!/usr/bin/env python
# coding: utf-8

import sys
import pybullet
from qibullet.camera import *
from qibullet.link import Link
from qibullet.joint import Joint

IS_VERSION_PYTHON_3 = sys.version_info[0] >= 3


class RobotVirtual:
    """
    Mother class representing a virtual robot
    """

    def __init__(self, description_file):
        """
        Constructor

        Parameters:
            description_file - The file giving the description of the virtual
            robot. For now, only URDF is handled
        """
        self.description_file = description_file
        self.physics_client = 0
        self.active_camera = None
        self.camera_dict = dict()
        self.joint_dict = dict()
        self.link_dict = dict()

    def loadRobot(self, translation, quaternion, physicsClientId=0):
        """
        Loads the robot into a simulation, loads the joints and the links
        descriptions. The joints are set to 0 rad.

        Parameters:
            translation - List containing 3 elements, the translation [x, y, z]
            of the robot in the WORLD frame
            quaternion - List containing 4 elements, the quaternion
            [x, y, z, q] of the robot in the WORLD frame
            physicsClientId - The id of the simulated instance in which the
            robot is supposed to be loaded

        Returns:
            boolean - True if the method ran correctly, False otherwise
        """
        try:
            self.physics_client = physicsClientId
            self.robot_model = pybullet.loadURDF(
                self.description_file,
                translation,
                quaternion,
                useFixedBase=False,
                globalScaling=1.0,
                physicsClientId=self.physics_client,
                flags=pybullet.URDF_USE_SELF_COLLISION |
                pybullet.URDF_USE_MATERIAL_COLORS_FROM_MTL)

        except pybullet.error as e:
            raise pybullet.error("Cannot load robot model: " + str(e))

        for i in range(pybullet.getNumJoints(
                self.robot_model,
                physicsClientId=self.physics_client)):
            if IS_VERSION_PYTHON_3:
                # PYTHON 3 version needs a conversion bytes to str
                joint_info = pybullet.getJointInfo(
                    self.robot_model,
                    i,
                    physicsClientId=self.physics_client)
                self.link_dict[joint_info[12].decode('utf-8')] =\
                    Link(joint_info)

                if joint_info[2] == pybullet.JOINT_PRISMATIC or\
                        joint_info[2] == pybullet.JOINT_REVOLUTE:
                    self.joint_dict[joint_info[1].decode('utf-8')] =\
                        Joint(joint_info)
            else:
                # PYTHON 2 Version
                joint_info = pybullet.getJointInfo(
                    self.robot_model,
                    i,
                    physicsClientId=self.physics_client)

                self.link_dict[joint_info[12]] = Link(joint_info)

                if joint_info[2] == pybullet.JOINT_PRISMATIC or\
                        joint_info[2] == pybullet.JOINT_REVOLUTE:
                    self.joint_dict[joint_info[1]] = Joint(joint_info)

    def getRobotModel(self):
        """
        Returns the pybullet model to which the module is associated.

        Returns:
            robot_model - The pybullet model of the robot
        """
        return self.robot_model

    def getPhysicsClientId(self):
        """
        Returns the id of the simulated instance in which the module is loaded.

        Returns:
            physics_client - The id of the simulation in which the robot
            (possessing the module) is spawned
        """
        return self.physics_client

    def setAngles(self, joint_names, joint_values, percentage_speeds):
        """
        Set angles on the robot's joints. Tests have to be performed by the
        child class to guarantee the validity of the input parameters.

        Parameters:
            joint_names - List of string containing the name of the joints
            to be controlled
            joint_values - List of values corresponding to the angles in
            radians to be applied
            percentage_speeds - Percentages of the max speed to be used for
            each joint, has to be strictly superior to 0 and inferior or equal
            to 1
        """
        try:
            assert len(joint_names) ==\
                len(joint_values) ==\
                len(percentage_speeds)

            assert all(
                speed >= 0.0 and speed <= 1.0 for speed in percentage_speeds)

        except AssertionError:
            raise pybullet.error("Error in the setAngles parameters")

        for joint_name, joint_value, percentage_speed in zip(
                joint_names,
                joint_values,
                percentage_speeds):

            joint_speed =\
                self.joint_dict[joint_name].getMaxVelocity() *\
                percentage_speed

            pybullet.setJointMotorControl2(
                self.robot_model,
                self.joint_dict[joint_name].getIndex(),
                pybullet.POSITION_CONTROL,
                targetPosition=joint_value,
                maxVelocity=joint_speed,
                force=self.joint_dict[joint_name].getMaxEffort(),
                physicsClientId=self.physics_client)

    def getAnglesPosition(self, joint_names):
        """
        Gets the position of the robot's joints in radians. If one of the joint
        doesn't exist, the method will raise a KeyError.

        Parameters:
            joint_names - List of string containing the names of the joints

        Returns:
            joint_positions - List of floats containing the joint's positions
        """
        joint_positions = list()

        for joint_name in joint_names:
            joint_positions.append(pybullet.getJointState(
                self.robot_model,
                self.joint_dict[joint_name].getIndex(),
                physicsClientId=self.physics_client)[0])

        return joint_positions

    def getAnglesVelocity(self, joint_names):
        """
        Gets the velocity of the robot's joints in rad/s. If one of the joint
        doesn't exist, the method will raise a KeyError.

        Parameters:
            joint_names - List of string containing the names of the joints

        Returns:
            joint_velocities - List of floats containing the joint's velocities
        """
        joint_velocities = list()

        for joint_name in joint_names:
            joint_velocities.append(pybullet.getJointState(
                self.robot_model,
                self.joint_dict[joint_name].getIndex(),
                physicsClientId=self.physics_client)[1])

        return joint_velocities

    def subscribeCamera(self, camera_id, resolution=Camera.K_QVGA):
        """
        Subscribe to the camera holding the camera id. WARNING: at the moment,
        only one camera can be subscribed.

        Parameters:
            camera_id - The id of the camera to be subscribed
            resolution - CameraResolution object, the resolution of the camera
        """
        try:
            self.active_camera = self.camera_dict[camera_id]
            self.active_camera.subscribe(resolution=resolution)

        except KeyError:
            print("This camera does not exist, use a valid camera id")

    def unsubscribeCamera(self, camera_id):
        """
        Unsubscribe from a camera, the one holding the camera id.

        Parameters:
            camera_id - The id of the camera to be unsubscribed
        """
        try:
            # If no active camera is found, nothing is unsubscribed
            assert self.active_camera is not None

            if self.active_camera.getCameraId() == camera_id:
                self.active_camera.unsubscribe()
                self.active_camera = None

        except KeyError:
            print("This camera does not exist, use a valid camera id")
        except AssertionError:
            pass

    def getCameraFrame(self):
        """
        Returns a camera frame. Be advised that the subscribeCamera method
        needs to be called beforehand, otherwise a pybullet error will be
        raised.

        Returns:
            frame - The current camera frame as a formatted numpy array,
            directly exploitable from OpenCV
        """
        try:
            assert self.active_camera is not None
            return self.active_camera.getFrame()

        except AssertionError:
            raise pybullet.error("No active camera, cannot retrieve any frame")

    def getCameraResolution(self):
        """
        Returns the resolution of the active camera. Be advised that the
        subscribeCamera method needs to be called beforehand, otherwise a
        pybullet error will be raised.

        Returns:
            resolution - a CameraResolution object describing the resolution of
            the active camera
        """
        try:
            assert self.active_camera is not None
            return self.active_camera.getResolution()

        except KeyError:
            raise pybullet.error("No active camera, resolution unavailable")

    def getCameraLink(self):
        """
        Returns the link of the active camera. Be advised that the
        subscribeCamera method needs to be called beforehand, otherwise a
        pybullet error will be raised.

        Returns:
            resolution - a Link object describing the link to which the active
            camera is attached
        """
        try:
            assert self.active_camera is not None
            return self.active_camera.getCameraLink()

        except KeyError:
            raise pybullet.error("No active camera, cannot retrieve any link")

    def getActiveCamera(self):
        """
        Returns the active camera of the robot.

        Returns:
            active_camera - Camera (CameraRgb or CameraDepth) object, the
            active camera of the robot. If there is no active camera, a None is
            returned
        """
        return self.active_camera

    def getPosition(self):
        """
        Gets the position of the robot's base in the world frame.

        Returns:
            x - The position of the robot's base on the x axis, in meters
            y - The positions of the robot's base on the y axis in meters
            theta - The rotation of the robot's base on the z axis in meters
        """
        position, quaternions = pybullet.getBasePositionAndOrientation(
            self.robot_model,
            physicsClientId=self.physics_client)

        theta = pybullet.getEulerFromQuaternion(quaternions)[2]
        return position[0], position[1], theta

    def isSelfColliding(self, link_names):
        """
        Specifies if a link is colliding with the rest of the virtual robot.

        Parameters:
            link_names - String or list of string containing the names of the
            links to be checked for self collision. WARNING: only the links
            with corresponding meshes should be used, otherwise the link cannot
            self collide

        Returns:
            self_colliding - Boolean, if True at least one of the links is self
            colliding
        """
        try:
            if type(link_names) is str:
                assert link_names in self.link_dict.keys()
                names = [link_names]
            else:
                assert set(link_names).issubset(self.link_dict.keys())
                names = list(link_names)

            for name in names:
                contact_tuple = pybullet.getContactPoints(
                    bodyA=self.robot_model,
                    bodyB=self.robot_model,
                    linkIndexA=self.link_dict[name].getIndex(),
                    physicsClientId=self.physics_client)
                contact_tuple += pybullet.getContactPoints(
                    bodyA=self.robot_model,
                    bodyB=self.robot_model,
                    linkIndexB=self.link_dict[name].getIndex(),
                    physicsClientId=self.physics_client)

                if len(contact_tuple) != 0:
                    return True

            return False

        except AssertionError:
            raise pybullet.error(
                "Unauthorized link checking for self collisions")

[((304, 32, 306, 48), 'pybullet.getBasePositionAndOrientation', 'pybullet.getBasePositionAndOrientation', (), '', False, 'import pybullet\n'), ((51, 31, 59, 59), 'pybullet.loadURDF', 'pybullet.loadURDF', (), '', False, 'import pybullet\n'), ((64, 23, 66, 52), 'pybullet.getNumJoints', 'pybullet.getNumJoints', (), '', False, 'import pybullet\n'), ((308, 16, 308, 60), 'pybullet.getEulerFromQuaternion', 'pybullet.getEulerFromQuaternion', ({(308, 48, 308, 59): 'quaternions'}, {}), '(quaternions)', False, 'import pybullet\n'), ((69, 29, 72, 56), 'pybullet.getJointInfo', 'pybullet.getJointInfo', (), '', False, 'import pybullet\n'), ((74, 20, 74, 36), 'qibullet.link.Link', 'Link', ({(74, 25, 74, 35): 'joint_info'}, {}), '(joint_info)', False, 'from qibullet.link import Link\n'), ((82, 29, 85, 56), 'pybullet.getJointInfo', 'pybullet.getJointInfo', (), '', False, 'import pybullet\n'), ((87, 49, 87, 65), 'qibullet.link.Link', 'Link', ({(87, 54, 87, 64): 'joint_info'}, {}), '(joint_info)', False, 'from qibullet.link import Link\n'), ((135, 18, 135, 69), 'pybullet.error', 'pybullet.error', ({(135, 33, 135, 68): '"""Error in the setAngles parameters"""'}, {}), "('Error in the setAngles parameters')", False, 'import pybullet\n'), ((248, 18, 248, 79), 'pybullet.error', 'pybullet.error', ({(248, 33, 248, 78): '"""No active camera, cannot retrieve any frame"""'}, {}), "('No active camera, cannot retrieve any frame')", False, 'import pybullet\n'), ((265, 18, 265, 76), 'pybullet.error', 'pybullet.error', ({(265, 33, 265, 75): '"""No active camera, resolution unavailable"""'}, {}), "('No active camera, resolution unavailable')", False, 'import pybullet\n'), ((282, 18, 282, 78), 'pybullet.error', 'pybullet.error', ({(282, 33, 282, 77): '"""No active camera, cannot retrieve any link"""'}, {}), "('No active camera, cannot retrieve any link')", False, 'import pybullet\n'), ((351, 18, 352, 65), 'pybullet.error', 'pybullet.error', ({(352, 16, 352, 64): '"""Unauthorized link checking for self collisions"""'}, {}), "('Unauthorized link checking for self collisions')", False, 'import pybullet\n'), ((79, 24, 79, 41), 'qibullet.joint.Joint', 'Joint', ({(79, 30, 79, 40): 'joint_info'}, {}), '(joint_info)', False, 'from qibullet.joint import Joint\n'), ((91, 53, 91, 70), 'qibullet.joint.Joint', 'Joint', ({(91, 59, 91, 69): 'joint_info'}, {}), '(joint_info)', False, 'from qibullet.joint import Joint\n')]