import numpy as np
import time
import copy as cp
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.interpolators import OTG

class DualDianaMed(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,
                 is_interpolate=False
                 ):
        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
        self.interpolator_left = None
        self.interpolator_right = None
        self.compute_qpos = np.array([1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16])
        if is_interpolate:
            self._initInterpolator(self.robot.arms)



        #camera view setting
        self.r_vis = None
        self.l_vis = None
        self.top = None
        self.angle = None
        self.front = None
        self.obs = None
        
        self.rew = None


    def actuate_J(self, q_target, qdot_target, Arm):
        """ Compute desired torque with robot dynamics modeling:
            > M(q)qdd + C(q, qd)qd + G(q) + tau_F(qd) = tau_ctrl + tau_env

        :param q_target: joint position
        :param qdot_target: joint velocity
        """
        # print(Arm.kdl_solver.NbOfJnt)
        acc_desire = [
            (self.robot.kp[i] * (q_target[i] - Arm.arm_qpos[i]) -
             (self.robot.kd[i] * Arm.arm_qvel[i])) for i in range(Arm.jnt_num)]
        qM = Arm.kdl_solver.getInertiaMat([Arm.arm_qpos[i] for i in range(len(Arm.joint_index))])
        tau_target = np.dot(qM, acc_desire) + np.array(
            [self.mj_data.joint(Arm.joint_index[i]).qfrc_bias[0] for i in range(len(Arm.joint_index))])
        # Send torque to simulation
        for i in range(7):
            self.mj_data.actuator(Arm.actuator_index[i]).ctrl = tau_target[i]    
        
        
    def reset(self):
        self._setInitPose()
        mj.mj_step(self.mj_model, self.mj_data)
        if self.interpolator_left is not None and self.interpolator_right is not None:
            self.interpolator_left.setOTGParam(self.robot.arms[0].arm_qpos,
                                        np.zeros(len(self.robot.arms[0].joint_index)))
            self.interpolator_right.setOTGParam(self.robot.arms[1].arm_qpos,
                                        np.zeros(len(self.robot.arms[1].joint_index)))
        


    def step(self,action):
        self.compute_qpos = action  #for observation !
        self.obs = self._get_obs()
        if self.interpolator_left is not None and self.interpolator_right is not None:
            self.interpolator_left.updateInput(action[:7], control_cycle=self.base_time)
            self.interpolator_right.updateInput(action[7:-2], control_cycle=self.base_time)

        ctrl_cur_time = time.time()
        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)
        self.base_time = time.time() - ctrl_cur_time
            

    def preStep(self, action):
        '''
        action size:
            left joints: 7 
            right joints: 7 
            left gripper left: 1
            right gripper left: 1
            total: 16
        '''
        if len(action) == 16:
            action_gripper_left = action[-2]
            action_gripper_right = action[-1]
            q_target_l = action[:7], 
            q_target_r= action[7:-2], 
            if self.interpolator_left is not None and self.interpolator_right is not None:
                try:
                    q_target_l, _ = self.interpolator_left.updateState()
                    q_target_r, _ = self.interpolator_right.updateState()
                except ValueError:
                    print(action)
            self.actuate_J(q_target_l, np.zeros(7), self.arm_left)
            self.actuate_J(q_target_r, np.zeros(7), self.arm_right)
            self.gripper_ctrl(action_gripper_left,self.arm_left)
            self.gripper_ctrl(action_gripper_right,self.arm_right)
        else:
            raise AttributeError

    def gripper_ctrl(self,ctrl,Arm):
        '''
        open:  100
        close: -100
        '''
        if -100 <= ctrl <= 100:
            self.mj_data.actuator(Arm.actuator_index[7]).ctrl = ctrl
        else:
            if ctrl>100:
                self.mj_data.actuator(Arm.actuator_index[7]).ctrl = 100
            else:
                self.mj_data.actuator(Arm.actuator_index[7]).ctrl = -100
            # raise ValueError("the torq should be given between -100 and 100 ")


    @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] = cp.deepcopy(self.mj_data.joint(self.arm_left.joint_index[i]).qpos[0])
        for i in range(self.arm_right.jnt_num):
            qpos_right[i] = cp.deepcopy(self.mj_data.joint(self.arm_right.joint_index[i]).qpos[0])
        return np.hstack([qpos_left,self.mj_data.actuator(self.arm_left.actuator_index[7]).ctrl,
                          qpos_right,self.mj_data.actuator(self.arm_right.actuator_index[7]).ctrl])


    def _get_obs(self):
        obs = collections.OrderedDict()
        obs['qpos'] = self.get_obs_qpos
        obs['action'] = self.compute_qpos
        obs['images'] = dict()
        obs['images']['top']   = self.top
        obs['images']['angle'] = self.angle
        obs['images']['r_vis'] = self.r_vis
        obs['images']['l_vis'] = self.l_vis
        obs['images']['front'] = self.front
        return obs

    def _get_image_obs(self):
        obs = collections.OrderedDict()
        obs['images'] = dict()
        obs['images']['top']   = self.top
        obs['images']['angle'] = self.angle
        obs['images']['r_vis'] = self.r_vis
        obs['images']['l_vis'] = self.l_vis
        obs['images']['front'] = self.front
        return obs
    
    def _get_qpos_obs(self):
        obs = collections.OrderedDict()
        obs['qpos'] = self.get_obs_qpos
        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 == 'r_vis':
            return self.r_vis
        elif self.cam == 'l_vis':
            return self.l_vis
        elif self.cam == 'front':
            return self.front
        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 not self.exit_flag:
            img_renderer.update_scene(self.mj_data,camera="rs_cam_right")
            self.r_vis = img_renderer.render()
            self.r_vis = self.r_vis[:, :, ::-1]
            img_renderer.update_scene(self.mj_data,camera="rs_cam_left")
            self.l_vis = img_renderer.render()
            self.l_vis = self.l_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]
            img_renderer.update_scene(self.mj_data,camera="front")
            self.front = img_renderer.render()
            self.front = self.front[:, :, ::-1]
            if self.cam_view is not None:
                cv2.imshow('Cam view', self.cam_view)
            cv2.waitKey(1)


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

    def _initInterpolator(self, Arm):
        
        self.interpolator_left = OTG(
            OTG_Dof=self.robot.left_arm.jnt_num,
            control_cycle=0.0005,
            max_velocity=0.05,
            max_acceleration=0.1,
            max_jerk=0.2
        )
        self.interpolator_left.setOTGParam(Arm[0].arm_qpos, Arm[0].arm_qvel)
        self.interpolator_right = OTG(
            OTG_Dof=self.robot.right_arm.jnt_num,
            control_cycle=0.0005,
            max_velocity=0.05,
            max_acceleration=0.1,
            max_jerk=0.2
        )
        self.interpolator_right.setOTGParam(Arm[1].arm_qpos, Arm[1].arm_qvel)

if __name__ == "__main__":
    from roboimi.assets.robots.diana_med import BiDianaMed

    env = DualDianaMed(
        robot=BiDianaMed(),
        is_render=True,
        control_freq=200,
        cam_view='top',
        is_interpolate=True
    )

    
    env.reset()
    
    for t in range(int(1e6)):
        action_l = np.array([0.0, 0.0, 0.0, 0.0, 0.0 ,0.0, 0.0])
        action_r = np.array([0.0, 0.0, 0.0, 0.0, 0.0 ,0.0, 0.0])
        gripper = np.array([-100.0,-100.0])
        action = np.hstack((action_l,action_r,gripper))
        env.step(action)
        l_q = np.zeros(7)
        l_q[3] = 0.0
        r_q = np.zeros(7)
        r_q[3] = 0.0
        print("\nkdl fk left =",env.arm_left.kdl_solver.getEEtf(l_q))
        # print("\n left pos=",env.robot.robot_data.body("left_link7").xpos,"\n")
        # print("\nleft xmat",env.robot.robot_data.body("left_link7").xmat,"\n")

        print("============================")
        print("\nkdl fk right =",env.arm_right.kdl_solver.getEEtf(r_q))
        # print("\n right pos=",env.robot.robot_data.body("right_link7").xpos,"\n")
        # print("\nright xmat",env.robot.robot_data.body("right_link7").xmat,"\n")

        # quat_left = np.zeros(4)
        # quat_right = np.zeros(4)
        # mj._functions.mju_mat2Quat(quat_left, rot_left)
        # print("quat_left =",quat_left,"\n")
        # mj._functions.mju_mat2Quat(quat_right, rot_right)
        # print("quat_right =",quat_right,"\n")
        if env.is_render:
            env.render()
        