import pybullet as p
import time
import pybullet_data
import numpy as np
import os
from pybullet_utils import bullet_client
import matplotlib.pyplot as plt
from Envs.pybullet.turtlebot.config import TurtleBotConfig
from Envs.pybullet.turtlebot.robot_locomotors import TurtleBot


class turtlebotStandAlone(object):
    """
    This class is a stand-alone software for interacting with the turtlebot
    """
    def __init__(self):
        self.config=TurtleBotConfig()

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

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

        self.turtleBotGUI()

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

    def build_arena(self):
        # introduce arena and the robot
        self.arena_wall = self._p.loadURDF(
            os.path.join(self.config.mediaPath, 'arena', 'arena_wall', 'arena_wall.urdf'),
            [0, 0, -0.02],
            [0.0, 0.0, 0.0, 1.0], flags=self._p.URDF_USE_MATERIAL_COLORS_FROM_MTL)
        self.arena_wall2 = self._p.loadURDF(
            os.path.join(self.config.mediaPath, 'arena', 'arena_wall', 'arena_wall2.urdf'),
            [0.25, 0, -0.02],
            [0.0, 0.0, 0.0, 1.0], flags=self._p.URDF_USE_MATERIAL_COLORS_FROM_MTL)

        texUid = p.loadTexture(
            os.path.join(self.config.mediaPath, 'arena', 'wallAppearance', 'texture', 'banded_0002.jpg'))
        p.changeVisualShape(self.arena_wall2, -1, textureUniqueId=texUid, rgbaColor=[1, 1, 1, 1])

        self.arena_floor = self._p.loadURDF(
            os.path.join(self.config.mediaPath, 'arena', 'arena_floor', 'arena_floor.urdf'),
            [0, 0, 0.03],
            [0.000000, 0.000000, 0.0, 1.0], flags=self._p.URDF_USE_MATERIAL_COLORS_FROM_MTL)

    def load_obj(self, objName):
        visualID = p.createVisualShape(shapeType=p.GEOM_MESH,
                                       fileName=os.path.join(self.config.mediaPath, 'objects', objName + '.obj'),
                                       visualFramePosition=[0, 0, 0],
                                       visualFrameOrientation=self._p.getQuaternionFromEuler([0, 0, 0]),

                                       rgbaColor=[1, 0, 0, 1],

                                       )

        collisionID = p.createCollisionShape(shapeType=p.GEOM_CYLINDER,
                                             radius=self.config.objectsRadius[
                                                        objName] + self.config.objectsExpandDistance[
                                                        objName] + self.config.robotCamOffset,
                                             height=0.3,
                                             collisionFramePosition=[0, 0, 0],
                                             collisionFrameOrientation=self._p.getQuaternionFromEuler([0, 0, 0]))

        self.blockUid = p.createMultiBody(baseMass=0,
                                          baseInertialFramePosition=[0, 0, 0],
                                          baseCollisionShapeIndex=collisionID,
                                          baseVisualShapeIndex=visualID,
                                          basePosition=[3, 3, 3],
                                          baseOrientation=self._p.getQuaternionFromEuler([0, 0, 0]))

    def turtleBotGUI(self):

        # flag and counters
        objName = 'cylinder'
        controlMethod ='rotPos'
        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))

        # build arena
        self.build_arena()

        # insert and place the robot
        self.robot = TurtleBot(config=self.config)
        self.robotID = self.robot.robot_ids
        self.robot.load_model()
        self.robot._p = self._p
        self.robot.eyes = self.robot.parts["eyes"]  # in the urdf, the link for the camera is called "eyes"

        # add an object into the scene
        blockx = self._p.addUserDebugParameter("blockx", -1.5, 1.5, 0.5)
        blocky = self._p.addUserDebugParameter("blocky", -1.5, 1.5, 0.5)
        blockz = self._p.addUserDebugParameter("blockz", -1.5, 1.5, self.robot.entityZ)
        self.load_obj(objName)

        for j in self.robot.ordered_joints:
            j.reset_joint_state(np.random.uniform(low=-0.1, high=0.1), 0)
        self.robot.robot_body.set_pose([0., 0., 0.06], self._p.getQuaternionFromEuler([0, 0, 0]))

        # control method
        actionBuf = [[0., 0.]] * 1  # it is used to simulate the delay happened in the actuator and communication
        if controlMethod=='rotPos':
            desiredTransVel = p.addUserDebugParameter("desired transitional velocity", -0.25, 0.25, 0)
            desiredRotPos = p.addUserDebugParameter("desired rotation position", - 2 *np.pi, 2* np.pi, 0)

        elif controlMethod=='point':
            desiredx = p.addUserDebugParameter("desired x", -0.95, 0.95, 0)
            desiredy = p.addUserDebugParameter("desired y", -0.95, 0.95, 0)
            desiredtheta = p.addUserDebugParameter("desired theta", -3.14, 3.14, 0)

        else:
            raise NotImplementedError

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

            # update obj position
            block_x = p.readUserDebugParameter(blockx)
            block_y = p.readUserDebugParameter(blocky)
            block_z = p.readUserDebugParameter(blockz)
            orn = self._p.getQuaternionFromEuler([0, -1.57, 0])
            self._p.resetBasePositionAndOrientation(self.blockUid, [block_x, block_y, block_z],
                                                    [orn[0], orn[1], orn[2], orn[3]])

            # update robot camera image
            if useRobotCam:
                self.robot.get_image()

            # control
            if controlMethod=='rotPos':
                # reading user input for control action
                desired_trans=p.readUserDebugParameter(desiredTransVel)
                desired_Rot_Pos = p.readUserDebugParameter(desiredRotPos)
                self.robot.desiredTransVel=desired_trans
                self.robot.desiredRotPos=desired_Rot_Pos
                self.robot.apply_action(action=[0,0], controlMethod=controlMethod)

            else:
                raise NotImplementedError

C=turtlebotStandAlone()
