import numpy as np
import time
from roboimi.envs.mujoco_base import MujocoEnv
import cv2
from roboimi.utils.interpolators import OTG
import mujoco as mj
import threading
from threading import Event
import collections
from roboimi.utils.constants import PUPPET_GRIPPER_POSITION_NORMALIZE_FN

class DualVX300S_Jnt_base(MujocoEnv):
    """

    :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 is_interpolate: Use interpolator while stepping.
    """

    def __init__(self,
                 robot=None,
                 is_render=False,
                 renderer="viewer",
                 control_freq=20,
                 cam_view=None
                 ):
        super().__init__(
            robot=robot,
            is_render=is_render,
            renderer=renderer,
            control_freq=control_freq
        )
        self.arm_left  = self.robot.left_arm
        self.arm_right = self.robot.right_arm
        self.base_time = self.control_timestep
        self.cam = cam_view
        #camera view setting
        self.vis = None
        self.top = None
        self.angle = None
        self.cam_is_open = False
        self.event = None
        # self.cam_start()

        
    def reset(self):
        # self.event = Event()
        # self.event.clear()
        # self.img_renderer = mj.Renderer(self.mj_model,height=480,width=640)
        self._setInitPose()
        mj.mj_step(self.mj_model, self.mj_data)
        

    def step(self,action=np.zeros(6)):
        for i in range(int(self.control_timestep / self.model_timestep)):
            if int(self.control_timestep / self.model_timestep) == 0:
                raise ValueError("Control frequency is too low. Checkout you are not in renderer mode."
                                 "Current Model-Timestep:{}".format(self.model_timestep))
            super().step(action)


    def preStep(self, action):
        if len(action) == self.robot.jnt_num + 4 and action is not None:
            action_l = action[:self.arm_left.jnt_num]
            action_r = action[self.arm_left.jnt_num:-4]
            self.move_J(action_l,self.arm_left)
            self.move_J(action_r,self.arm_right)
            grasp_ctrl = action[-4:]
            self.gripper_ctrl(grasp_ctrl)
        else:
            raise AttributeError
        

    def move_J(self, jnt : np.ndarray, Arm):
        if len(jnt) == Arm.jnt_num:
            for i in range(Arm.jnt_num):
                self.mj_data.actuator(Arm.actuator_index[i]).ctrl = jnt[i]
        else:
            raise AttributeError
        

    def gripper_ctrl(self, ctrl : np.ndarray):
        if len(ctrl) == 4:
            self.mj_data.actuator(self.arm_left.actuator_index[6]).ctrl = ctrl[0]
            self.mj_data.actuator(self.arm_left.actuator_index[7]).ctrl = ctrl[1]
            self.mj_data.actuator(self.arm_right.actuator_index[6]).ctrl = ctrl[2]
            self.mj_data.actuator(self.arm_right.actuator_index[7]).ctrl = ctrl[3]
        else:
            raise AttributeError
        

    @property
    def get_obs_qpos(self):
        qpos_left = np.zeros(self.arm_left.jnt_num)
        qpos_right = np.zeros(self.arm_right.jnt_num)
        for i in range(self.arm_left.jnt_num):
            qpos_left[i] = self.mj_data.joint(self.arm_left.joint_index[i]).qpos
        for i in range(self.arm_right.jnt_num):
            qpos_right[i] = self.mj_data.joint(self.arm_right.joint_index[i]).qpos
        left_gripper_qpos = self.mj_data.joint(self.arm_left.gripper_index[0]).qpos
        right_gripper_qpos = self.mj_data.joint(self.arm_right.gripper_index[0]).qpos
        return np.hstack([qpos_left,left_gripper_qpos,qpos_right,right_gripper_qpos])
    

    def _get_obs(self):
        obs = collections.OrderedDict()
        obs['qpos'] = self.get_obs_qpos
        obs['qvel'] = self.mj_data.qvel
        obs['env_state'] = self.get_env_state()
        obs['images'] = dict()
        obs['images']['top']   = self.top
        obs['images']['angle'] = self.angle
        obs['images']['vis']   = self.vis
        return obs


    def get_env_state(self):
        raise NotImplementedError

    
    def _get_reward(self):
        raise NotImplementedError

    @property
    def cam_view(self):
        if self.cam == 'top':
            return self.top
        elif self.cam == 'angle':
            return self.angle
        elif self.cam == 'vis':
            return self.vis
        else:
            raise AttributeError("please input right name")

    def camera_viewer_test(self):
            self.img_renderer = mj.Renderer(self.mj_model,height=480,width=640)
            self.img_renderer.update_scene(self.mj_data,camera="front_close")
            self.vis = self.img_renderer.render()
            self.vis = self.vis[:, :, ::-1]
            self.img_renderer.update_scene(self.mj_data,camera="top")
            self.top = self.img_renderer.render()
            self.top =self.top[:, :, ::-1]
            self.img_renderer.update_scene(self.mj_data,camera="angle")
            self.angle = self.img_renderer.render()
            self.angle = self.angle[:, :, ::-1]
            

    def camera_viewer(self):
        img_renderer = mj.Renderer(self.mj_model,height=480,width=640)
        print("delay ========",time.time()-self.mj_time)
        # cv2.namedWindow('Cam view',cv2.WINDOW_NORMAL)
        while True:
            # print("init win")
            # if self.event.isSet():
            #     cv2.destroyAllWindows()
            #     print("event is set")
            #     break
            
            img_renderer.update_scene(self.mj_data,camera="front_close")
            self.vis = img_renderer.render()
            self.vis = self.vis[:, :, ::-1]
            img_renderer.update_scene(self.mj_data,camera="top")
            self.top = img_renderer.render()
            self.top =self.top[:, :, ::-1]
            img_renderer.update_scene(self.mj_data,camera="angle")
            self.angle = img_renderer.render()
            self.angle = self.angle[:, :, ::-1]
            
            
            # cv2.imshow('Cam view', self.top)
            # cv2.waitKey(1)
            



    def cam_start(self):
        
        self.cam_thread = threading.Thread(target=self.camera_viewer)
        self.cam_thread.daemon = True
        self.cam_thread.start()
        






if __name__ == "__main__":
    from roboimi.assets.robots.vx300s import VX300s_insert_jnt,VX300s_transfer_jnt

    env = DualVX300S_Jnt_base(
        robot=VX300s_insert_jnt(),
        is_render=True,
        control_freq=200,
        cam_view='top'
    )

    # env = DualVX300S_EE_base(
    #     robot=VX300s_transfer(),
    #     is_render=True,
    #     control_freq=200,
    #     cam_view='top'
    # )

    # env = transfer
    env.reset()
    # env.cam_start()
    action_l = np.zeros(6)
    action_l[0]=0.9
    action_r = np.zeros(6)
    action_r[0] = -0.0
    gripper_ctrl = np.array([0.0,-0,0.0,0.0])
    action = np.hstack((action_l,action_r,gripper_ctrl))
    # action = np.array([1,2,3,4,5,6,7,8,9,10,11,12])
    for t in range(int(1e6)):
        env.step(action)
        if env.is_render:
            env.render()
        