import pybullet as p
import time
import pybullet_data
import numpy as np
import os
from pybullet_utils import bullet_client
from Envs.pybullet.kuka.RSI2.config import KukaConfig
from Envs.pybullet.kuka.robot_manipulators import Kuka


class KukaStandAlone(object):
    """
    This class is a stand-alone software for interacting with the Kuka iiwa
    """
    def __init__(self):
        self.config=KukaConfig()
        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))

        objects = self._p.loadSDF(os.path.join(pybullet_data.getDataPath(), "kuka_iiwa/kuka_with_gripper2.sdf"))
        self.robot=Kuka(config=self.config)
        self.robot.robot_ids = objects[0]
        self.robot._p = self._p

        # anchor the robot
        self._p.resetBasePositionAndOrientation(self.robot.robot_ids, [-0.100000, 0.000000, 0.070000],
                                                [0.000000, 0.000000, 0.000000, 1.000000])

        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
        rp=[] # restposes 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]))
            rp.append(0)
            jd.append(jointInfo[6])
        print('joint lower limits', ll)
        print('joint upper limits', ul)
        print('joint ranges', jr)
        print('joint damping', jd)


        # load table and obj
        self.tableUid = self._p.loadURDF(os.path.join(pybullet_data.getDataPath(), "table/table.urdf"),
                                         [0.5, 0.0, self.config.tableZ],
                                         [0.0, 0.0, 0.0, 1.0])
        objPath = os.path.join(self.config.mediaPath, 'objects', 'key.urdf')

        # randomize the locations of the blocks
        orn = self._p.getQuaternionFromEuler([0, 0, 0])  # do not randomize the orientation
        randomx = np.random.uniform(self.config.xMin + 0.05, self.config.xMax - 0.05)  # 0.5~0.7  fix: 0.55
        randomy = np.random.uniform(self.config.yMin + 0.05, self.config.yMax - 0.45)  # -0.2~-0.1 fix: -0.15

        for i in range(len(self.config.objList)):
            self.objUidList.append(self._p.loadURDF(objPath))
            self._p.resetBasePositionAndOrientation(self.objUidList[i],
                                                    [randomx, randomy + i * self.config.objInterval,
                                                     self.config.objZ],
                                                    orn)

        gripperBase = self._p.addUserDebugParameter("gripperBase", ll[7], ul[7], 0)
        leftFinger = self._p.addUserDebugParameter("leftFinger", ll[8], ul[8], 0)
        leftFingerTip = self._p.addUserDebugParameter("leftFingerTip", ll[10], ul[10], 0)
        rightFinger = self._p.addUserDebugParameter("rightFinger", ll[11], ul[11], 0)
        rightFingerTip = self._p.addUserDebugParameter("rightFingerTip", ll[13], ul[13], 0)

        if controlMethod == 'joint':
            J1 = self._p.addUserDebugParameter("J1", ll[0], ul[0], 0)
            J2 = self._p.addUserDebugParameter("J2", ll[1], ul[1], 0)
            J3 = self._p.addUserDebugParameter("J3", ll[2], ul[2], 0)
            J4 = self._p.addUserDebugParameter("J4", ll[3], ul[3], 0)
            J5 = self._p.addUserDebugParameter("J5", ll[4], ul[4], 0)
            J6 = self._p.addUserDebugParameter("J6", ll[5], ul[5], 0)
            J7 = self._p.addUserDebugParameter("J7", ll[6], ul[6], 0)
        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.22, 0.7, 0.22)
            self.robot.robot_specific_reset(0.3,0,0.22)

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

            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)

                gb = self._p.readUserDebugParameter(gripperBase)
                lf = self._p.readUserDebugParameter(leftFinger)
                lft = self._p.readUserDebugParameter(leftFingerTip)
                rf = self._p.readUserDebugParameter(rightFinger)
                rft = self._p.readUserDebugParameter(rightFingerTip)

                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, 7, self._p.POSITION_CONTROL, targetPosition=gb,
                                              force=self.config.positionControlMaxForce)
                self._p.setJointMotorControl2(self.robot.robot_ids, 8, self._p.POSITION_CONTROL, targetPosition=lf,
                                              force=self.config.fingerAForce)
                self._p.setJointMotorControl2(self.robot.robot_ids, 11, self._p.POSITION_CONTROL, targetPosition=lft,
                                              force=self.config.fingerBForce)

                self._p.setJointMotorControl2(self.robot.robot_ids, 10, self._p.POSITION_CONTROL, targetPosition=rf,
                                              force=self.config.fingerTipForce)
                self._p.setJointMotorControl2(self.robot.robot_ids, 13, self._p.POSITION_CONTROL, targetPosition=rft,
                                              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)

            if useRobotCam:
                self.robot.get_image()

C=KukaStandAlone()
