import numpy as np
import time, math
from roboimi.envs.mujoco_base import MujocoEnv
import cv2
from roboimi.utils.interpolators import OTG
import mujoco as mj
import threading
from scipy.spatial.transform import Rotation as RR


class DualArmCamAruEnv(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=200,
                 is_interpolate=False
                 ):
        super().__init__(
            robot=robot,
            is_render=is_render,
            renderer=renderer,
            control_freq=control_freq
        )
        self.arm = self.robot.left_arm
        self.aruco_arm = self.robot.aruco_arm
        self.base_time = self.control_timestep
        self.nv = self.mj_model.nv
        if is_interpolate:
            self._initInterpolator(self.arm)
            self._initInterpolator(self.aruco_arm)
        self.marker_size = 0.05
        self.aruco_dict = 'DICT_4X4_100'
        self.distCoeffs = np.zeros(5)
        self.aruco_parameters = cv2.aruco.DetectorParameters()
        self.cameraMatrix = self.fov2Intrinsic()
        
        self.Detected_already = False
        self.image = None

    def action_size(self):
        return self.robot.jnt_num

    def fov2Intrinsic(self,fovy=55.0, width=320, height=240):
        aspect = width * 1.0 / height
        fovx = math.degrees(2 * math.atan(aspect * math.tan(math.radians(fovy / 2))))

        cx = width / 2.0
        cy = height / 2.0
        fx = cx / math.tan(fovx * math.pi / 180 * 0.5)
        fy = cy / math.tan(fovy * math.pi / 180 * 0.5)

        K = [[fx, 0, cx],
             [0, fy, cy],
             [0,  0,  1]]
        return np.array(K, dtype=float)
    
    def camera_viewer(self):
        img_renderer = mj.Renderer(self.mj_model)
        cv2.namedWindow('Cam view',cv2.WINDOW_NORMAL)
        while True:
            img_renderer.update_scene(self.mj_data, camera="0_cam")
            org = img_renderer.render()
            self.image = org[:, :, ::-1]
            cv2.imshow('Cam view', self.image)
            cv2.waitKey(1)

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


    def actuate_base(self, base):
        for i in range(len(self.robot.base_joint_index)):
            acc_desire = 100 * (base - self.mj_data.joint(self.robot.base_joint_index[i]).qpos[0]) - \
                         66 * (self.mj_data.joint(self.robot.base_joint_index[i]).qvel[0])
            tau_target = np.subtract(np.dot(2, acc_desire),
                                     self.mj_data.joint(self.robot.base_joint_index[i]).qfrc_bias[0])
            self.mj_data.actuator(self.robot.base_actuator_index[i]).ctrl = tau_target

    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
        """
        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(len(Arm.joint_index))]
        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(len(Arm.actuator_index)):
            self.mj_data.actuator(Arm.actuator_index[i]).ctrl = tau_target[i]
    
    def aruco_detection(self):
        try:
            if self.aruco_dict != "DICT_4X4_100":
                raise ValueError("Invalid aruco dictionary ID: " + self.aruco_dict)
        except ValueError as e:
            print(str(e))
        dictionary_id = cv2.aruco.__getattribute__(self.aruco_dict)
        aruco_dictionary = cv2.aruco.getPredefinedDictionary(dictionary_id)
        R_vec = np.array([0.0, 0.0, 0.0])
        T_vec = np.array([0.0, 0.0, 0.0])
        if self.image is not None:
            corners, marker_ids, _ = cv2.aruco.detectMarkers(self.image,
                                                             aruco_dictionary, self.cameraMatrix, self.distCoeffs,
                                                             parameters=self.aruco_parameters)
            if marker_ids is not None:
                self.Detected_already = True
                if cv2.__version__ > '4.0.0':
                    rvecs, tvecs, _ = cv2.aruco.estimatePoseSingleMarkers(corners,
                                                                          self.marker_size, self.cameraMatrix,
                                                                          self.distCoeffs)
                else:
                    rvecs, tvecs = cv2.aruco.estimatePoseSingleMarkers(corners,
                                                                       self.marker_size, self.cameraMatrix,
                                                                       self.distCoeffs)
                T_vec = np.array([tvecs[0][0][0], tvecs[0][0][1], tvecs[0][0][2]])
                R_vec = np.array([rvecs[0][0][0], rvecs[0][0][1], rvecs[0][0][2]])
            else:
                self.Detected_already = False

        return R_vec, T_vec
    

    def mj_aruco_detection(self):
        try:
            if self.aruco_dict != "DICT_4X4_100":
                raise ValueError("Invalid aruco dictionary ID: " + self.aruco_dict)
        except ValueError as e:
            print(str(e))
        dictionary_id = cv2.aruco.__getattribute__(self.aruco_dict)
        aruco_dictionary = cv2.aruco.getPredefinedDictionary(dictionary_id)
        rvec = np.zeros(3)
        tvec = np.zeros(3)
        # print("llll")
        if self.image is not None:
            corners, marker_ids, _ = cv2.aruco.detectMarkers(self.image,
                                                             aruco_dictionary, self.cameraMatrix, self.distCoeffs,
                                                             parameters=self.aruco_parameters)
            
            if marker_ids is not None:
                self.Detected_already = True
                '''
                create aruco Transformation 
                w : world frame of mujoco
                a : aruco frame 
                c : camera frame 
                '''
                t_wa = self.getBodyPos("aruco")
                R_wa = self.getBodyMat("aruco")
                T_wa = np.eye(4)
                T_wa[:3,:3] = R_wa
                T_wa[:3,3] = t_wa

                t_wc = self.getBodyPos("0_cam")
                R_wc = self.getBodyMat("0_cam")
                T_wc = np.eye(4)
                T_wc[:3,:3] = R_wc
                T_wc[:3,3] = t_wc
                T_cw = np.linalg.inv(T_wc)
                T_ca = T_cw@T_wa
                R_ca = T_ca[:3,:3]
                t_ca = T_ca[:3,3]
                rvec = RR.from_matrix(R_ca).as_rotvec()
                tvec = t_ca

            else:
                self.Detected_already = False
                # rvec = np.zeros(3)
                # tvec = np.zeros(3)

        return rvec, tvec, self.Detected_already



    def preStep(self, action):
        q_target_l = action[1:8]
        qdot_target_l =  np.zeros(self.arm.jnt_num)
        q_target_r, qdot_target_r = action[8:15], np.zeros(self.aruco_arm.jnt_num)

        if self.arm.interpolator and self.aruco_arm.interpolator is not None:
            q_target_l, qdot_target_l = self.arm.interpolator.updateState()
            q_target_r, qdot_target_r = self.aruco_arm.interpolator.updateState()

        self.actuate_J(q_target_l, qdot_target_l, self.arm)
        self.actuate_J(q_target_r, qdot_target_r, self.aruco_arm)
        
        
    def step(self, action):
        '''
        action : base  1
                 left  7
                 right 7
        '''
        if self.arm.interpolator and self.aruco_arm.interpolator is not None:
            self.arm.interpolator.updateInput(action[1:8], control_cycle=self.base_time)
            self.aruco_arm.interpolator.updateInput(action[8:15], 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 reset(self):
        self._setInitPose()
        mj.mj_step(self.mj_model, self.mj_data)
        

    def _initInterpolator(self, Arm):
        Arm.interpolator = OTG(
            OTG_Dof=7,
            control_cycle=0.0005,
            max_velocity=0.05,
            max_acceleration=0.1,
            max_jerk=0.2
        )
        Arm.interpolator.setOTGParam(Arm.init_pose, np.zeros(Arm.jnt_num))

        

# if __name__ == "__main__":
#     from roboimi.assets.robots.bidiana_med_vs import BiDianaMedVS

#     env = DualArmCamAruEnv(
#         robot=BiDianaMedVS(),
#         is_render=True,
#         control_freq=200,
#         is_interpolate=True
#     )

#     env.reset()
#     env.cam_start()
#     action_base = np.array([0.0])
#     # action_L = np.array([0.0, -0.4943951023931953, 0.0, 1.853981633974483, 0.0, 0.7853981633974483, 0.0]) #basic pos
#     # action_R = np.array([0.0, -0.1943951023931953, 0.0, 2.153981633974483, 0.0, 0.7853981633974483, 0.0]) #basic pos
#     action_L = np.zeros(7)
#     action_R = np.zeros(7)
#     action_L[3] = 1.57
#     action_R[3] = 1.57
#     for t in range(int(1e6)):
#         action_L[3] = action_L[3]  + 0.0000001*t 
#         action = np.concatenate((action_base, action_L, action_R), axis=0)
#         # print("pos=",env.mj_data.body("link7").xpos,"\n")
#         # print("kdl solve p ==",env.arm.kdl_solver.getEEtf(action_L).p,"\n")
#         # print("mat=\n",env.mj_data.body("link7").xmat[0:3],"\n")
#         # print("mat=\n",env.mj_data.body("link7").xmat[3:6],"\n")
#         # print("mat=\n",env.mj_data.body("link7").xmat[6:9],"\n")
#         # print("kdl solve M ==\n",env.arm.kdl_solver.getEEtf(action_L).M,"\n")
#         # print("===================================")
#         print(action_L)
#         env.step(action)
#         if env.is_render:
#             env.render()

