import numpy as np
import time
from roboimi.envs.vx300s_jnt_base import DualVX300S_Jnt_base
import cv2
from roboimi.utils.interpolators import OTG
import mujoco as mj
import threading
from pyquaternion import Quaternion
import copy



class DualVX300S_Jnt_Insert(DualVX300S_Jnt_base):
    """

    :param is_render: Choose if use the renderer to render the scene or not.
    :param renderer: Choose official renderer with "viewer",
            another renderer with "mujoco_viewer"
    :param control_freq: Upper-layer control frequency.
            Note that high frequency will cause high time-lag.
    :param cam_view: Choice of Camera view.
    """

    def __init__(self,
                 robot=None,
                 is_render=False,
                 renderer="viewer",
                 control_freq=200,
                 cam_view=None
                 ):
        super().__init__(
            robot=robot,
            is_render=is_render,
            renderer=renderer,
            control_freq=control_freq,
            cam_view=cam_view
        )
        self.obs = None
        self.rew = None
        self.max_reward = 4
        self.box_pose = None
        self.num = 0
        
        self.cam_start()
        


    def step(self,action=np.zeros(14)):
        self.obs = self._get_obs()
        self.rew = self._get_reward()
        super().step(action)


    def reset(self,box_pose):
        self.mj_data.joint('red_peg_joint').qpos[0] = box_pose[0]
        self.mj_data.joint('red_peg_joint').qpos[1] = box_pose[1]
        self.mj_data.joint('red_peg_joint').qpos[2] = box_pose[2]
        self.mj_data.joint('red_peg_joint').qpos[3] = 1.0
        self.mj_data.joint('red_peg_joint').qpos[4] = 0.0
        self.mj_data.joint('red_peg_joint').qpos[5] = 0.0
        self.mj_data.joint('red_peg_joint').qpos[6] = 0.0

        self.mj_data.joint('blue_socket_joint').qpos[0] = box_pose[3]
        self.mj_data.joint('blue_socket_joint').qpos[1] = box_pose[4]
        self.mj_data.joint('blue_socket_joint').qpos[2] = box_pose[5]
        self.mj_data.joint('blue_socket_joint').qpos[3] = 1.0
        self.mj_data.joint('blue_socket_joint').qpos[4] = 0.0
        self.mj_data.joint('blue_socket_joint').qpos[5] = 0.0
        self.mj_data.joint('blue_socket_joint').qpos[6] = 0.0
        super().reset()
       
        self.cam_flage = True
        t=0
        while self.cam_flage:
            if(type(self.top)==type(None)):
                time.sleep(0.01)
                # self.obs = self._get_obs()
                t+=1
            else:
                self.cam_flage=False
                print("\nsleep time=",0.01*t)
        self.obs = self._get_obs()


    def preStep(self, action=np.zeros(14)):
        action_left = action[:self.arm_left.jnt_num]
        action_right = action[self.arm_right.jnt_num+1:-1]
        if self.num >= 15000:
            action_grip = np.array([action[self.arm_left.jnt_num]*0.6,-action[self.arm_left.jnt_num]*0.6,action[-1]*0.4,-action[-1]*0.4])
            # if self.num >=27000:
            #     action_grip = np.array([action[self.arm_left.jnt_num]*0.4,-action[self.arm_left.jnt_num]*0.4,action[-1]*0.3,-action[-1]*0.3])

        else:
            action_grip = np.array([action[self.arm_left.jnt_num],-action[self.arm_left.jnt_num],action[-1],-action[-1]])
        action_jnt_base = np.hstack((action_left,action_right,action_grip))
        super().preStep(action_jnt_base)
        self.num += 1
        

        
         
    def get_env_state(self):
        return self.mj_data.qpos.copy()[16:]
    
    def _get_reward(self):
        all_contact_pairs = []
        for collision_num in range(self.mj_data.ncon):
            geom1 = self.mj_data.contact[collision_num].geom1
            geom2 = self.mj_data.contact[collision_num].geom2
            name_geom_1 = self.getID2Name('geom',geom1)
            name_geom_2 = self.getID2Name('geom',geom2)
            contact_pair = (name_geom_1, name_geom_2)
            all_contact_pairs.append(contact_pair)

        touch_right_gripper = ("red_peg", "vx300s_right/10_right_gripper_finger") in all_contact_pairs
        touch_left_gripper = ("socket-1", "vx300s_left/10_left_gripper_finger") in all_contact_pairs or \
                             ("socket-2", "vx300s_left/10_left_gripper_finger") in all_contact_pairs or \
                             ("socket-3", "vx300s_left/10_left_gripper_finger") in all_contact_pairs or \
                             ("socket-4", "vx300s_left/10_left_gripper_finger") in all_contact_pairs

        peg_touch_table = ("red_peg", "table") in all_contact_pairs
        socket_touch_table = ("socket-1", "table") in all_contact_pairs or \
                             ("socket-2", "table") in all_contact_pairs or \
                             ("socket-3", "table") in all_contact_pairs or \
                             ("socket-4", "table") in all_contact_pairs
        peg_touch_socket = ("red_peg", "socket-1") in all_contact_pairs or \
                           ("red_peg", "socket-2") in all_contact_pairs or \
                           ("red_peg", "socket-3") in all_contact_pairs or \
                           ("red_peg", "socket-4") in all_contact_pairs
        pin_touched = ("red_peg", "pin") in all_contact_pairs

        reward = 0
        if touch_left_gripper and touch_right_gripper: # touch both
            reward = 1
        if touch_left_gripper and touch_right_gripper and (not peg_touch_table) and (not socket_touch_table): # grasp both
            reward = 2
        if peg_touch_socket and (not peg_touch_table) and (not socket_touch_table): # peg and socket touching
            reward = 3
        if pin_touched: # successful insertion
            reward = 4
        return reward
        



# class DualVX300S_Jnt_Transfer(DualVX300S_Jnt_base):
#     """

#     :param is_render: Choose if use the renderer to render the scene or not.
#     :param renderer: Choose official renderer with "viewer",
#             another renderer with "mujoco_viewer"
#     :param control_freq: Upper-layer control frequency.
#             Note that high frequency will cause high time-lag.
#     :param cam_view: Choice of Camera view.
#     """

#     def __init__(self,
#                  robot=None,
#                  is_render=False,
#                  renderer="viewer",
#                  control_freq=200,
#                  cam_view=None
#                  ):
#         super().__init__(
#             robot=robot,
#             is_render=is_render,
#             renderer=renderer,
#             control_freq=control_freq,
#             cam_view=cam_view
#         )
#         self.obs = None
#         self.rew = None
#         self.max_reward = 4
        
#     def step(self,action=np.zeros(14)):
#         self.obs = self._get_obs()
#         self.rew = self._get_reward()
#         super().step(action)
        


#     def reset(self):
#         super().reset()

#     def preStep(self, action):
#         pass

#     def get_env_state(self):
#         return self.mj_data.qpos.copy()[16:]

#     def _get_reward(self):
#         all_contact_pairs = []
#         for collision_num in range(self.mj_data.ncon):
#             geom1 = self.mj_data.contact[collision_num].geom1
#             geom2 = self.mj_data.contact[collision_num].geom2
#             name_geom_1 = self.getID2Name('geom',geom1)
#             name_geom_2 = self.getID2Name('geom',geom2)
#             contact_pair = (name_geom_1, name_geom_2)
#             all_contact_pairs.append(contact_pair)

#         touch_left_gripper = ("red_box", "vx300s_left/10_left_gripper_finger") in all_contact_pairs
#         touch_right_gripper = ("red_box", "vx300s_right/10_right_gripper_finger") in all_contact_pairs
#         touch_table = ("red_box", "table") in all_contact_pairs

#         reward = 0
#         if touch_right_gripper:
#             reward = 1
#         if touch_right_gripper and not touch_table: # lifted
#             reward = 2
#         if touch_left_gripper: # attempted transfer
#             reward = 3
#         if touch_left_gripper and not touch_table: # successful transfer
#             reward = 4
#         return reward
        

def make_sim_env(task_name):
    # if 'sim_transfer_cube' in task_name:
    #     from roboimi.assets.robots.vx300s import VX300s_transfer_jnt
    #     env = DualVX300S_Jnt_Transfer(
    #         robot=VX300s_transfer_jnt(),
    #         is_render=True,
    #         control_freq=30,
    #         cam_view='top'
    #     )
    if 'sim_insertion' in task_name:
        from roboimi.assets.robots.vx300s import VX300s_insert_jnt
        env = DualVX300S_Jnt_Insert(
            robot=VX300s_insert_jnt(),
            is_render=True,
            control_freq=30,
            cam_view='top'
        )
    else:
        raise NotImplementedError
    return env


if __name__ == "__main__":
    task_name = 'sim_insertion'
    env = make_sim_env(task_name)
    box = np.array([0.0,0.5,1.05,-0.2,-0.8,0.05])
    env.reset(box)
    action_l = np.array([-0.0,-0.0,-0.5,
                         0.0,0.0,0.0])
    action_r = np.array([-0.5,0.0,0.0,
                         0.0,0.0,0.0])
    action_grip = np.array([-0.1,-0.1])
    action = np.hstack((action_l,action_r,action_grip))
    ti=0
    for t in range(int(1e6)):
        env.step(action)
        ti +=1
        print("ti==",ti)
        if env.is_render:
            env.render()
        # if time == 1000:
        #     env.reset()
        #     time= 0