import pybullet as p
import time
import pybullet_data
import numpy as np
import os
from pybullet_utils import bullet_client
from Envs.pybullet.kinova_gen3.config import KinovaGen3Config
from Envs.pybullet.kinova_gen3.robot_manipulators import KinovaGen3


class KukaStandAlone(object):
    """
    This class is a stand-alone software for interacting with the Kuka iiwa
    """

    def __init__(self):
        self.config = KinovaGen3Config()
        self.robot = None
        self.tableUid = None
        self.objUidList = []

        # GUI
        self._p = bullet_client.BulletClient(connection_mode=p.GUI)

        p.setAdditionalSearchPath(pybullet_data.getDataPath())  # optionally
        p.setGravity(0, 0, -9.8)

        self.KukaGUI()

    def __del__(self):
        self._p.disconnect()

    def KukaGUI(self):

        # flag and counters
        controlMethod = 'position'
        stepcounter = 0
        frameSkip = 6
        useRobotCam = True

        self.config.render = True
        self.config.mediaPath = 'media'

        # setup simulation
        assert self._p.isNumpyEnabled() == 1
        self._p.setPhysicsEngineParameter(fixedTimeStep=1. / 240. * frameSkip, numSolverIterations=50,
                                          numSubSteps=(frameSkip - 1))

        self.robot = KinovaGen3(config=self.config)
        self.robot.load_model()
        self.robot._p = self._p

        # anchor the robot
        self._p.resetBasePositionAndOrientation(self.robot.robot_ids, [0.0, 0.0, 0.07], [0.0, 0.0, 0.0, 1.0])

        self.robot.numJoints = self._p.getNumJoints(self.robot.robot_ids)

        ll = []  # lower limits for inverse kinematics null space
        ul = []  # upper limits for inverse kinematics null space
        jr = []  # joint ranges for inverse kinematics null space
        jd = []  # joint damping coefficents for inverse kinematics null space

        for i in range(self.robot.numJoints):
            jointInfo = p.getJointInfo(self.robot.robot_ids, i)
            print('jointIndex: ', jointInfo[0], 'jointName: ', jointInfo[1], 'jointType: ', jointInfo[2], 'linkName: ', jointInfo[12])
            ll.append(jointInfo[8])
            ul.append(jointInfo[9])
            jr.append(abs(jointInfo[8]) + abs(jointInfo[9]))
            jd.append(jointInfo[6])
        print('joint lower limits', ll)
        print('joint upper limits', ul)
        print('joint ranges', jr)
        print('joint damping', jd)

        # after we get all the join info we set it to the number of non-fixed joint
        self.robot.numJoints = len(self.robot.jdict)

        # load table and obj
        self.tableUid = self._p.loadURDF(os.path.join(self.config.mediaPath, 'table', 'table.urdf'),
                                         [0.66, -0.14, self.config.tableZ],
                                         [0.0, 0.0, 0.0, 1.0])
        objPath = os.path.join(self.config.mediaPath, 'objects', 'key.urdf')
        for i in range(len(self.config.objList)):
            self.objUidList.append(self._p.loadURDF(objPath, useFixedBase=1))

        # add a base for the robot
        visualID = self._p.createVisualShape(shapeType=self._p.GEOM_CYLINDER,
                                             radius=0.08,
                                             length=0.02,
                                             visualFramePosition=[0, 0, 0],
                                             visualFrameOrientation=self._p.getQuaternionFromEuler([0, 0, 0]),
                                             rgbaColor=[0, 0, 0, 1]
                                             )

        baseID = self._p.createMultiBody(baseMass=0,
                                        baseInertialFramePosition=[0, 0, 0],
                                        baseVisualShapeIndex=visualID,
                                        basePosition=[0,0,-0.01],
                                        baseOrientation=self._p.getQuaternionFromEuler([0, 0, 0]))

        # a wall behind  the robot for domain randomization
        wall_visualID = self._p.createVisualShape(shapeType=self._p.GEOM_BOX,
                                                  halfExtents=[0.05, 1.3, 0.7],
                                                  visualFramePosition=[0, 0, 0],
                                                  visualFrameOrientation=self._p.getQuaternionFromEuler([0, 0, 0]),
                                                  rgbaColor=[0.6, 0.5, 0.2, 1]
                                                  )

        wallID = self._p.createMultiBody(baseMass=0,
                                         baseInertialFramePosition=[0, 0, 0],
                                         baseVisualShapeIndex=wall_visualID,
                                         basePosition=[-0.3, -0.2, -0.3],
                                         baseOrientation=self._p.getQuaternionFromEuler([0, 0, 0]))

        # randomize the locations of the blocks
        orn = self._p.getQuaternionFromEuler([0, 0, 0])  # do not randomize the orientation
        blockx = self._p.addUserDebugParameter("blockx", self.config.xMin + 0.05, self.config.xMax - 0.05, 0)
        blocky = self._p.addUserDebugParameter("blocky", self.config.yMin + 0.05, self.config.yMax - 0.45, 0.4)

        finger_joint = self._p.addUserDebugParameter("finger_joint", -1.0, 1.0, 0.723)
        left_inner_finger_joint = self._p.addUserDebugParameter("left_inner_finger_joint", -1.0, 1.0, 0.723)
        left_inner_knuckle_joint = self._p.addUserDebugParameter("left_inner_knuckle_joint", -1.0, 1.0, -0.723)
        right_outer_knuckle_joint = self._p.addUserDebugParameter("right_outer_knuckle_joint", -1.0, 1.0, -0.723)
        right_inner_finger_joint = self._p.addUserDebugParameter("right_inner_finger_joint", -1.0, 1.0, 0.723)
        right_inner_knuckle_joint = self._p.addUserDebugParameter("right_inner_knuckle_joint", -1.0, 1.0, -0.723)

        if controlMethod == 'joint':
            J1 = self._p.addUserDebugParameter("J1", ll[0], ul[0], 0)
            J2 = self._p.addUserDebugParameter("J2", ll[1], ul[1], 0.4)
            J3 = self._p.addUserDebugParameter("J3", ll[2], ul[2], 0)
            J4 = self._p.addUserDebugParameter("J4", ll[3], ul[3], 1.57)
            J5 = self._p.addUserDebugParameter("J5", ll[4], ul[4], 0)
            J6 = self._p.addUserDebugParameter("J6", ll[5], ul[5], 1.1)
            J7 = self._p.addUserDebugParameter("J7", ll[6], ul[6], 1.57)
        elif controlMethod == 'position':
            X = self._p.addUserDebugParameter("x", 0.3, 1, 0.3)
            Y = self._p.addUserDebugParameter("y", -0.5, 0.5, 0)
            Z = self._p.addUserDebugParameter("z", 0.15, 0.4, self.config.endEffectorHeight)
            self.robot.robot_specific_reset(0.3, 0, self.config.endEffectorHeight)

        workSpaceDebugLine=[]
        # draw debug line
        start = [self.config.xMin, self.config.yMin, self.config.objZ]
        end = [self.config.xMax, self.config.yMin, self.config.objZ]
        workSpaceDebugLine.append(p.addUserDebugLine(start, end, self.config.rayMissColor, lineWidth=5))
        start = [self.config.xMin, self.config.yMax, self.config.objZ]
        end = [self.config.xMax, self.config.yMax, self.config.objZ]
        workSpaceDebugLine.append(p.addUserDebugLine(start, end, self.config.rayMissColor, lineWidth=5))
        start = [self.config.xMax, self.config.yMin, self.config.objZ]
        end = [self.config.xMax, self.config.yMax, self.config.objZ]
        workSpaceDebugLine.append(p.addUserDebugLine(start, end, self.config.rayMissColor, lineWidth=5))
        start = [self.config.xMin, self.config.yMin, self.config.objZ]
        end = [self.config.xMin, self.config.yMax, self.config.objZ]
        workSpaceDebugLine.append(p.addUserDebugLine(start, end, self.config.rayMissColor, lineWidth=5))

        # start simulation
        while 1:
            # step simulation
            stepcounter = stepcounter + 1
            p.stepSimulation()
            time.sleep(1. / 240. * frameSkip)  # camera frame rate

            bx = self._p.readUserDebugParameter(blockx)
            by = self._p.readUserDebugParameter(blocky)
            for i in range(len(self.config.objList)):

                self._p.resetBasePositionAndOrientation(self.objUidList[i],
                                                        [bx, by + i * self.config.objInterval,
                                                         self.config.objZ],
                                                        orn)

            if controlMethod == 'joint':
                joint1 = self._p.readUserDebugParameter(J1)
                joint2 = self._p.readUserDebugParameter(J2)
                joint3 = self._p.readUserDebugParameter(J3)
                joint4 = self._p.readUserDebugParameter(J4)
                joint5 = self._p.readUserDebugParameter(J5)
                joint6 = self._p.readUserDebugParameter(J6)
                joint7 = self._p.readUserDebugParameter(J7)

                fj = self._p.readUserDebugParameter(finger_joint)
                li = self._p.readUserDebugParameter(left_inner_finger_joint)
                lik = self._p.readUserDebugParameter(left_inner_knuckle_joint)
                rok = self._p.readUserDebugParameter(right_outer_knuckle_joint)
                rif = self._p.readUserDebugParameter(right_inner_finger_joint)
                rik = self._p.readUserDebugParameter(right_inner_knuckle_joint)

                jointPositions = [joint1, joint2, joint3, joint4, joint5, joint6, joint7]
                for i in range(self.config.endEffectorIndex + 1):
                    self._p.setJointMotorControl2(bodyUniqueId=self.robot.robot_ids, jointIndex=i,
                                                  controlMode=self._p.POSITION_CONTROL,
                                                  targetPosition=jointPositions[i], targetVelocity=0,
                                                  force=self.config.positionControlMaxForce,
                                                  positionGain=self.config.positionControlPositionGain,
                                                  velocityGain=self.config.positionControlVelGain)

                self._p.setJointMotorControl2(self.robot.robot_ids, 13, self._p.POSITION_CONTROL, targetPosition=fj,
                                              force=self.config.positionControlMaxForce)
                self._p.setJointMotorControl2(self.robot.robot_ids, 15, self._p.POSITION_CONTROL, targetPosition=li,
                                              force=self.config.fingerAForce)
                self._p.setJointMotorControl2(self.robot.robot_ids, 17, self._p.POSITION_CONTROL, targetPosition=lik,
                                              force=self.config.fingerBForce)

                self._p.setJointMotorControl2(self.robot.robot_ids, 18, self._p.POSITION_CONTROL, targetPosition=rok,
                                              force=self.config.fingerTipForce)
                self._p.setJointMotorControl2(self.robot.robot_ids, 20, self._p.POSITION_CONTROL, targetPosition=rif,
                                              force=self.config.fingerTipForce)
                self._p.setJointMotorControl2(self.robot.robot_ids, 22, self._p.POSITION_CONTROL, targetPosition=rik,
                                              force=self.config.fingerTipForce)

            elif controlMethod == 'position':
                x = self._p.readUserDebugParameter(X)
                y = self._p.readUserDebugParameter(Y)
                z = self._p.readUserDebugParameter(Z)

                self.robot.desiredEndEffectorPos[0] = x
                self.robot.desiredEndEffectorPos[1] = y
                self.robot.desiredEndEffectorPos[2] = z
                self.robot.applyAction(eeCommands=[0., 0., 0.], controlMethod=controlMethod)

            self.robot.ray_test(self.objUidList)
            if useRobotCam:
                self.robot.get_image(externalCamEyePosition=self.config.externalCamEyePosition,
                                     externalCamTargetPosition=self.config.externalCamTargetPosition)

            print('eeLinkState', p.getLinkState(self.robot.robot_ids, self.config.endEffectorIndex)[0])

C = KukaStandAlone()
