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
import collections
import time
from roboimi.utils.constants import PUPPET_GRIPPER_POSITION_CLOSE
from roboimi.utils.constants import PUPPET_GRIPPER_POSITION_NORMALIZE_FN,PUPPET_GRIPPER_VELOCITY_NORMALIZE_FN
CLOSE_GRIPPER = np.array([
                            -0.0409,
                            0.0409,
                            PUPPET_GRIPPER_POSITION_CLOSE,
                            -PUPPET_GRIPPER_POSITION_CLOSE,
                        ])


class DualVX300S_EE_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
        # np.copyto(physics.data.ctrl, close_gripper_control)

        #camera view setting
        self.vis = None
        self.top = None
        self.angle = None
        # self.cam_start()
        
        
    def reset(self):
        self._setInitPose()
        mj.mj_step(self.mj_model, self.mj_data)
        # super().reset()


    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) == 18:
            action_gripper = action[-4:]
            action_cart = action[:14]
            pos_l = action_cart[:3]
            quat_l = action_cart[3:7]
            pos_r = action_cart[7:10]
            quat_r = action_cart[10:14]
            self.motion_Cap(pos_l,quat_l,self.arm_left)
            self.motion_Cap(pos_r,quat_r,self.arm_right)
            self.gripper_ctrl(action_gripper)
        else:
            raise AttributeError
        
        
    def motion_Cap(self,pos : np.ndarray, quat : np.ndarray, Arm):
        self.mj_data.mocap_pos[Arm.ee_mocap] = pos
        self.mj_data.mocap_quat[Arm.ee_mocap] = quat


    def gripper_ctrl(self,ctrl=CLOSE_GRIPPER):
        if len(ctrl)==4:
            self.mj_data.actuator(self.arm_left.actuator_index[0]).ctrl = ctrl[0]
            self.mj_data.actuator(self.arm_left.actuator_index[1]).ctrl = ctrl[1]
            self.mj_data.actuator(self.arm_right.actuator_index[0]).ctrl = ctrl[2]
            self.mj_data.actuator(self.arm_right.actuator_index[1]).ctrl = ctrl[3]
        else:
            raise AttributeError("length of gripper ctrl vector should be 4 ")


    @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])


    @property
    def get_obs_qvel(self):
        qvel_left = np.zeros(self.arm_left.jnt_num)
        qvel_right = np.zeros(self.arm_right.jnt_num)
        for i in range(self.arm_left.jnt_num):
            qvel_left[i] = self.mj_data.joint(self.arm_left.joint_index[i]).qvel
        for i in range(self.arm_right.jnt_num):
            qvel_right[i] = self.mj_data.joint(self.arm_right.joint_index[i]).qvel
        left_gripper_qvel = PUPPET_GRIPPER_VELOCITY_NORMALIZE_FN(self.mj_data.joint(self.arm_left.gripper_index[0]).qvel)
        right_gripper_qvel = PUPPET_GRIPPER_VELOCITY_NORMALIZE_FN(self.mj_data.joint(self.arm_right.gripper_index[0]).qvel)
        return np.hstack([qvel_left,left_gripper_qvel,qvel_right,right_gripper_qvel])


    def _get_obs(self):
        obs = collections.OrderedDict()
        obs['qpos'] = self.get_obs_qpos
        obs['qvel'] = self.get_obs_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
        #used in scripted policy to obtain starting pose
        obs['mocap_pose_left'] = np.concatenate([self.mj_data.mocap_pos[0], self.mj_data.mocap_quat[0]]).copy()
        obs['mocap_pose_right'] = np.concatenate([self.mj_data.mocap_pos[1], self.mj_data.mocap_quat[1]]).copy()
        obs['gripper_ctrl'] = self.mj_data.ctrl.copy()
        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(self):
        img_renderer = mj.Renderer(self.mj_model,height=480,width=640)
        cv2.namedWindow('Cam view',cv2.WINDOW_NORMAL)
        while True:
            img_renderer.update_scene(self.mj_data,camera="front_close")
            self.vis = img_renderer.render()
            img_renderer.update_scene(self.mj_data,camera="top")
            self.top = img_renderer.render()
            img_renderer.update_scene(self.mj_data,camera="angle")
            self.angle = img_renderer.render()
            # org = img_renderer.render()
            # self.image = org[:, :, ::-1]
            cv2.imshow('Cam view', self.cam_view)
            cv2.waitKey(1)


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



# if __name__ == "__main__":
#     from roboimi.assets.robots.vx300s import VX300s_insert_ee,VX300s_transfer_ee

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

    

#     env.reset()
#     action_l = np.array([-0.2,0.3,0.3,
#                          1.0,0.0,0.0,0.0])
#     # action_r = np.array([0.3,0.1,0.3,
#     #                      1.0,0.0,0.0,0.0])
#     action_r = np.array([0.2,0.3,0.3,
#                          1.0,0.0,0.0,0.0])
#     gripper = np.array([0.0,0.0,0.0,0.0])
#     action = np.hstack((action_l,action_r,gripper))
    
#     for t in range(int(1e6)):
#         env.step(action)
#         if env.is_render:
#             env.render()
        