import numpy as np
import time
from roboimi.envs.vx300s_ee_base import DualVX300S_EE_base
import cv2
from roboimi.utils.interpolators import OTG
import mujoco as mj
import threading
from pyquaternion import Quaternion
from roboimi.utils.constants import PUPPET_GRIPPER_POSITION_NORMALIZE_FN
from roboimi.utils.constants import PUPPET_GRIPPER_POSITION_UNNORMALIZE_FN
from roboimi.utils.constants import PUPPET_GRIPPER_POSITION_CLOSE 
from roboimi.utils.act_ex_utils import sample_insertion_pose
import copy as cp



class DualVX300S_EE_Insert(DualVX300S_EE_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.num = 0
        
    def step(self,action=np.zeros(16)):
        self.obs = self._get_obs()
        self.rew = self._get_reward()
        super().step(action)

    def reset(self):
        peg_pose, socket_pose = sample_insertion_pose()
        self.mj_data.joint('red_peg_joint').qpos[0]=peg_pose[0]
        self.mj_data.joint('red_peg_joint').qpos[1]=peg_pose[1]
        self.mj_data.joint('red_peg_joint').qpos[2]=peg_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] = socket_pose[0]
        self.mj_data.joint('blue_socket_joint').qpos[1] = socket_pose[1]
        self.mj_data.joint('blue_socket_joint').qpos[2] = socket_pose[2]
        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()

    def preStep(self, action):
        if isinstance(action,np.ndarray) and len(action)==16:
            act_grip = action[-2:]
            act_cart = action[:-2]
            g_left_ctrl = PUPPET_GRIPPER_POSITION_UNNORMALIZE_FN(act_grip[0])
            g_right_ctrl = PUPPET_GRIPPER_POSITION_UNNORMALIZE_FN(act_grip[1])
            
            act_grip = np.array([g_left_ctrl, -g_left_ctrl, g_right_ctrl, -g_right_ctrl])
            action = np.hstack((act_cart,act_grip))
            super().preStep(action)
            self.num += 1
        else:
            raise AttributeError("size should be 16")
         
    def get_env_state(self):
        box_pose = np.zeros(6)
        for i in range(3):
            box_pose[i] = cp.deepcopy(self.mj_data.joint('red_peg_joint').qpos[i])
            box_pose[3+i] = cp.deepcopy(self.mj_data.joint('blue_socket_joint').qpos[i])
        return box_pose
    
    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_EE_Transfer(DualVX300S_EE_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_ee_sim_env(task_name):
    if 'sim_transfer_cube' in task_name:
        from roboimi.assets.robots.vx300s import VX300s_transfer_ee
        env = DualVX300S_EE_Transfer(
            robot=VX300s_transfer_ee(),
            is_render=True,
            control_freq=200,
            cam_view='top'
        )
    elif 'sim_insertion' in task_name:
        from roboimi.assets.robots.vx300s import VX300s_insert_ee
        env = DualVX300S_EE_Insert(
            robot=VX300s_insert_ee(),
            is_render=False,
            control_freq=30,
            cam_view='top'
        )
    else:
        raise NotImplementedError
    return env


if __name__ == "__main__":
    task_name = 'sim_insertion'
    env = make_ee_sim_env(task_name)
    env.reset()
    action_l = np.array([-0.0,0.1,0.3,
                         1.0,0.0,0.0,0.0])
    action_r = np.array([0.1,-0.1,0.1,
                         1.0,0.0,0.0,0.0])
    action_gripper = np.array([1.0,-0.3])
    action = np.hstack((action_l,action_r,action_gripper))
    time=0
    for t in range(int(100)):
        print("main step ", t)
        env.step(action)
        print("obs=",env.obs['gripper_ctrl'])
        time +=1
        if env.is_render:
            env.render()
        