import gym
from gym import error, spaces, utils
from gym.utils import seeding
import time
import os
import pybullet as p
import pybullet_data
import math
import numpy as np
import random
from math import sqrt


class JumpFrogEnv(gym.Env):
    metadata = {'render.modes': ['human']}

    def __init__(self,mode = p.GUI):
        self.action_space_limit = [0.3,0.3,0.2,0.3,0.3,0.2]
        p.connect(mode)
        p.resetDebugVisualizerCamera(cameraDistance=1.5, cameraYaw=0, \
                                     cameraPitch=-40, cameraTargetPosition=[0.55, -0.35, 0.2])
        # 6 motors joints to controll
        upper_l = np.array([0.6, 0.2, -0.9, 0.6, 0.2, -0.9])
        lower_l = np.array([0.1,  -0.2, -1.1, 0.1,  -0.2, -1.1])
        action_area = (upper_l - lower_l) *self.action_space_limit
        print(action_area)
        self.action_space = spaces.Box(lower_l, upper_l)
        self.observation_space = spaces.Box(np.array([-4]*20),np.array([4]*20))
        self.joint_index_to_control = [9, 8, 10, 2, 1, 3]

    def switchmode(self,mode = p.GUI):
        p.disconnect()
        p.connect(mode)
    def reward(self,state_robot):
        pose_reward = -sqrt((state_robot[3])**2)*2 - sqrt(state_robot[4]**2)*2 - sqrt(state_robot[5]**2)*2 - state_robot[0]**2 - state_robot[1]**2
        print("pose_reward")
        print(pose_reward)
        if state_robot[2] > 0.18:
            height_reward = 2
        else:
            height_reward = 0
        print("height_reward")
        print(height_reward)
        print("Time_reward")
        time_reward = sqrt((time.time() - self.agent_life)*10)
        print(time_reward)
        reward = pose_reward + height_reward
        for i in range(3):
            if (self.action[i] - self.action[i+3]) <0.1:
                reward += 1
        if state_robot[2] > 0.28:
            reward -= 10
        for i in range(6):
            if sqrt((self.action[i] - self.prev_action[i])**2) > self.action_space_limit[i]:
                reward -= 1
                print("exceed_limit")
        return reward
    def robot_state(self):
        body = p.getLinkState(self.pandaUid, 1,computeLinkVelocity= 1)
        body_state = list(body[0] + p.getEulerFromQuaternion(body[1])) + list(body[6]) + list(body[7])
        motor_state = []
        for i in range(6):
            motor_state.append(p.getJointState(self.pandaUid,self.joint_index_to_control[i])[0])
        r_f_angle = p.getEulerFromQuaternion(p.getLinkState(self.pandaUid, 5)[1])[0]
        l_f_angle = p.getEulerFromQuaternion(p.getLinkState(self.pandaUid, 12)[1])[0]
        foot_state = [r_f_angle,l_f_angle]
        state = body_state + motor_state + foot_state
        print("STATE")
        print(state)
        return state
    def step(self, action):
        self.action = action
        print("action:")
        print(action)
        p.configureDebugVisualizer(p.COV_ENABLE_SINGLE_STEP_RENDERING)
        for i in range(len(self.joint_index_to_control)):
            p.setJointMotorControl2(int(self.pandaUid),int(self.joint_index_to_control[i]),p.POSITION_CONTROL,targetPosition = action[i])
        for i in range(14):
            if i not in self.joint_index_to_control:
                p.setJointMotorControl2(int(self.pandaUid), i, p.POSITION_CONTROL,0)
        p.setJointMotorControl2(int(self.pandaUid), 7, p.POSITION_CONTROL, 1)
        p.setJointMotorControl2(int(self.pandaUid), 14, p.POSITION_CONTROL, 1)
        p.stepSimulation()
        state_robot = self.robot_state()
        r_f_angle = p.getEulerFromQuaternion(p.getLinkState(self.pandaUid, 5)[1])[0]
        l_f_angle = p.getEulerFromQuaternion(p.getLinkState(self.pandaUid, 12)[1])[0]

        reward = self.reward(state_robot)


        if state_robot[2] < 0.18:
            reward -= 10
        #     done = True
        # else:
        #     done = False
        if self.sim_timestep < 0:
            done = True
        else:
            done = False
        self.sim_timestep -=1
        print("reward: ")
        print(reward)
        info = {"is_success":None}
        observation = state_robot
        self.prev_action = action
        return observation, reward, done, info

    def reset(self):
        self.sim_timestep = 400
        self.agent_life = time.time()
        p.resetSimulation()
        p.configureDebugVisualizer(p.COV_ENABLE_RENDERING, 0)
        p.setGravity(0, 0, -10)

        p.setAdditionalSearchPath(pybullet_data.getDataPath())
        planeUid = p.loadURDF("plane.urdf")
        rest_poses = [0] * 15
        cubeStartPos = [0, 0, 0.28]
        cubeStartOrientation = p.getQuaternionFromEuler([0, 0, 0])
        self.pandaUid = p.loadURDF("C:/Users\Lorne\PycharmProjects\Simulation\jumpfrog.urdf", cubeStartPos,cubeStartOrientation)

        # c = p.createConstraint(self.pandaUid,3,self.pandaUid,6,p.JOINT_POINT2POINT,[1,0,0],[0, 0, -0.03],[0, 0, -0.04])
        p.setJointMotorControl2(self.pandaUid, 7, p.POSITION_CONTROL, 2)
        p.setJointMotorControl2(self.pandaUid, 14, p.POSITION_CONTROL, 2)
        p.setJointMotorControl2(self.pandaUid, 3, p.POSITION_CONTROL, -1.3)
        p.setJointMotorControl2(self.pandaUid, 10, p.POSITION_CONTROL, -1.3)
        p.setJointMotorControl2(self.pandaUid, 2, p.POSITION_CONTROL, 0.3)
        p.setJointMotorControl2(self.pandaUid, 9, p.POSITION_CONTROL, 0.3)
        p.setJointMotorControl2(self.pandaUid, 6, p.POSITION_CONTROL, -0.5)
        p.setJointMotorControl2(self.pandaUid, 13, p.POSITION_CONTROL, -0.5)
        self.prev_action = [0.3, 0, -1.3, 0.3, 0, -1.3]
        for i in range(70):
            p.stepSimulation()
        observation = self.robot_state()
        p.configureDebugVisualizer(p.COV_ENABLE_RENDERING, 1)
        return np.array(observation)


    def render(self, mode='human', width=500, height=500):
        pass

    def close(self):
        p.disconnect()