import numpy as np
from roboimi.assets.robots.bidiana_med_vs import BiDianaMedVS
from roboimi.envs.visual_servo.double_cam_aruco_base import DualArmCamAruEnv
import roboimi.utils.KDL_utils.transform as T
import math
from scipy.spatial.transform import Rotation as RR




class DualVisualServoEnv(DualArmCamAruEnv):
    """
    :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,
            is_interpolate=is_interpolate
        )
        self.t_ec = np.array([0.065054, 0.00068038, -0.013835]) #from XML (note: ec = description of cam pos vector based on end-effector)
        self.quat_ec = np.array([-2.59734e-06, 0.707105, 0.707108, -2.59735e-06])
        self.R_ec = RR.from_quat([0.707105, 0.707108, -2.59735e-06, -2.59734e-06]).as_matrix() #from XML (note: w x y z in XML but x y z w in RR)
        self.arm_base_pos,self.arm_base_mat = self.arm.CartPoseComp()
        self.aruco_base_pos,self.aruco_base_mat = self.aruco_arm.CartPoseComp()
        self.rvec = None
        self.tvec = None
        self.rv = None
        self.tv = None
        self.dt = 0.01
        self.prev_err = np.zeros(6)
        self.detect_buf = []
        self.is_detect = True
        # self.object_to_desire_camera = np.array([[ 0.999612, -0.010182,  0.02593, -0.022], 
        #                                          [-0.009818, -1.000338, -0.01413,  0.002],
        #                                          [  0.02607,   0.01387, -1.00005,  0.167],
        #                                          [        0,         0,        0,      1]])
        # self.kp = 1
        
        
        # rvec = np.array([-2.16881568 , 2.16116001 , 0.09872025])
        rvec = np.array([0.0,0.0,1.5708])
        # tvec = np.array([0.00480954, -0.02087546,  0.23213283])
        tvec = np.array([0.0,0.0,0.3])
        # tvec = np.array([0.00480954, -0.12087546,  0.23213283])
        
        rmat = RR.from_rotvec(rvec).as_matrix()
        self.T_cam_aruco = np.eye(4)
        self.T_cam_aruco[:3,:3] = rmat
        self.T_cam_aruco[:3,3] = tvec
        T_aruco_cam = np.linalg.inv(self.T_cam_aruco)
        self.object_to_desire_camera = T_aruco_cam
        self.T_acd = T_aruco_cam
        

    def getEEVel(self,Arm):
        current_qpos = Arm.arm_qpos
        current_qvel = Arm.arm_qvel
        current_q_states = Arm.kdl_solver.setJntArrayVel(current_qpos,current_qvel)
        twist = Arm.kdl_solver.fkSolverVel(current_q_states)
        return twist

    def Vel2qpos(self,cart_vel,Arm):
        '''
        frame world : w
        frame cam   : c
        cart vel    : V_wc (v_wc, w_wc) 
        '''
        _,mat_bias = Arm.get_Arm_bias_frame() #R_wb
        mat_bias = np.array([[mat_bias[0],mat_bias[1],mat_bias[2]],
                             [mat_bias[3],mat_bias[4],mat_bias[5]],
                             [mat_bias[6],mat_bias[7],mat_bias[8]]])
        mat_bias = np.linalg.inv(mat_bias) #R_bw
        current_qpos = Arm.arm_qpos
        cart_vel[:3] = mat_bias@cart_vel[:3] #v_bc = R_bw@v_wc
        cart_vel[3:] = mat_bias@cart_vel[3:] #w_bc = R_bw@w_wc
        qvel = Arm.kdl_solver.diffikSolver(current_qpos,cart_vel)
        qpos = qvel*self.dt + current_qpos
        return qpos

    def ik_solve(self,pos_goal : np.ndarray, quat_goal : np.ndarray, Arm):
        pos_bias,mat_bias = Arm.get_Arm_bias_frame()
        pos_bias = np.array([pos_bias[0],pos_bias[1],pos_bias[2]])
        mat_bias = np.array([[mat_bias[0],mat_bias[1],mat_bias[2]],
                             [mat_bias[3],mat_bias[4],mat_bias[5]],
                             [mat_bias[6],mat_bias[7],mat_bias[8]]])
        mat_bias = np.linalg.inv(mat_bias)
        mat_goal = T.quat2Mat(quat_goal)
        p_goal = mat_bias@(pos_goal-pos_bias)
        mat_goal = mat_bias@mat_goal
        return Arm.kdl_solver.ikSolver(p_goal, mat_goal, Arm.arm_qpos)

    def error_derivation(self,current_err):
        dot = np.array([0.0,0.0,0.0,0.0,0.0,0.0])
        for i in range(6):
            dot[i] = (current_err[i] - self.prev_err[i])/self.dt
        return dot[:3],dot[3:]
    # def pbvs_controller(self, aruco_to_cam_r, aruco_to_cam_p):
    #     t_acd = self.desire_p
    #     # print(aruco_to_cam_r)
    #     t_ca = aruco_to_cam_p
    #     R_ca = RR.from_rotvec(aruco_to_cam_r).as_matrix()
    #     t_ac = np.linalg.inv(R_ca)@t_ca.reshape(3, 1)
    #     R_acd = RR.from_quat(self.desire_r).as_matrix()
    #     R_ccd = R_ca@R_acd
    #     thu = RR.from_matrix(R_ccd).as_rotvec()
    #     T_ca_s = np.array([[       0,  -t_ca[2],   t_ca[1]],
    #                        [ t_ca[2],         0,  -t_ca[0]],
    #                        [-t_ca[1],   t_ca[0],         0]])
    #     w_cam = thu.reshape(3, 1)
    #     w_cam = np.zeros(3).reshape(3,1)
    #     v_cam = -(R_ca@(t_acd - t_ac)+ T_ca_s@w_cam)

    #     error = np.concatenate([v_cam,w_cam],axis=0)
    #     R_bc = self.getBodyMat("0_end_effector")@self.R_ec
    #     # R_bc = np.eye(3)
    #     Rv = np.vstack((np.hstack((R_bc, np.zeros((3, 3)))),
    #                     np.hstack((np.zeros((3, 3)), R_bc))))
    #     V_camera = np.matmul(Rv, error)
    #     return V_camera
    

    def pbvs_controller_origin(self, gain, kp, aruco_to_cam_r, aruco_to_cam_p):
        # gain = 1E-02
        var_lambda = kp
        t_ca = aruco_to_cam_p
        r_ca = aruco_to_cam_r
        R_ca = RR.from_rotvec(r_ca).as_matrix()
        T_ca = np.eye(4)
        T_ca[:3,:3] = R_ca
        T_ca[:3,3] = t_ca
        T_ccd = T_ca@self.T_acd
        R_ccd = T_ccd[:3,:3]
        T_cdc = np.linalg.inv(T_ccd) #formulation needed
        
        t_cdc = T_cdc[:3,3]
        R_cdc = T_cdc[:3,:3]
        thu = RR.from_matrix(R_cdc).as_rotvec()
        err_v = R_ccd@t_cdc
        error = np.concatenate([err_v,thu])
        if np.linalg.norm(error) > gain:
            Lin_v = -var_lambda*err_v
            Ang_v = -var_lambda*thu
            Vcam = np.concatenate([Lin_v,Ang_v])
            R_wc = self.getBodyMat("0_cam")
            Rv = np.vstack((np.hstack((R_wc, np.zeros((3, 3)))),
                            np.hstack((np.zeros((3, 3)), R_wc))))
            output = np.matmul(Rv,Vcam)
            print("cam vel=",output)
            return output
        else:
            print("err don't exist ")
            return np.zeros(6)
        
    def pbvs_controller_PD(self, gain, kp, kd, aruco_to_cam_r, aruco_to_cam_p):
        var_lambda = kp
        t_ca = aruco_to_cam_p
        r_ca = aruco_to_cam_r
        R_ca = RR.from_rotvec(r_ca).as_matrix()
        T_ca = np.eye(4)
        T_ca[:3,:3] = R_ca
        T_ca[:3,3] = t_ca
        T_ccd = T_ca@self.T_acd
        R_ccd = T_ccd[:3,:3]
        T_cdc = np.linalg.inv(T_ccd) #formulation needed
        
        t_cdc = T_cdc[:3,3]
        R_cdc = T_cdc[:3,:3]
        # print("here")
        thu = RR.from_matrix(R_cdc).as_rotvec()
        err_t = R_ccd@t_cdc
        error = np.concatenate([err_t,thu])
        err_v,err_w = self.error_derivation(error)
        if np.linalg.norm(error) > gain:
            Lin_v = -(var_lambda*err_t + kd*err_v)
            Ang_v = -(var_lambda*thu + kd*err_w)
            self.prev_err = error
            Vcam = np.concatenate([Lin_v,Ang_v])
            R_wc = self.getBodyMat("0_cam")
            Rv = np.vstack((np.hstack((R_wc, np.zeros((3, 3)))),
                            np.hstack((np.zeros((3, 3)), R_wc))))
            output = np.matmul(Rv,Vcam)
            print("cam vel=",output)
            return output
        else:
            print("err don't exist ")
            return np.zeros(6)

    # def quat_to_axis_angle(self,q):
    #     a, b, c, d = q
    #     v = np.array([a, b, c])
    #     angle = 2 * np.arccos(d)
    #     if angle == 0:
    #         axis = np.array([1, 0, 0])
    #     else:
    #         axis = v / np.linalg.norm(v)
    #         axis *= np.sin(angle / 2)
    #     return axis, angle

    # def pbvs_controller_PD(self,aruco_to_cam_r, aruco_to_cam_p):
    #     t_acd = self.object_to_desire_camera[:3,3]
    #     # R_acd = RR.from_quat(self.desire_quat).as_matrix()
    #     R_acd = self.object_to_desire_camera[:3,:3]
    #     T_acd = np.eye(4)
    #     T_acd[:3,:3] = R_acd
    #     T_acd[:3,3] = t_acd.reshape(1,3)
    #     t_ca = aruco_to_cam_p
    #     R_ca = RR.from_rotvec(aruco_to_cam_r).as_matrix()
    #     T_ca = np.eye(4)
    #     T_ca[:3,:3] = R_ca
    #     T_ca[:3,3] = t_ca
    #     T_ac = np.linalg.inv(T_ca)
    #     T_ccd = T_ca@T_acd

    #     R_ccd = T_ccd[:3,:3] # for test
    #     T_cdc = np.linalg.inv(T_ccd)
    #     R_cdc = T_cdc[:3,:3]
    #     t_cdc = T_cdc[:3,3]
        
    #     gain = 1E-04
    #     t_ac = T_ac[:3,3]
    #     t_ac_s = np.array([[       0,  -t_ac[2],   t_ac[1]],
    #                        [ t_ac[2],         0,  -t_ac[0]],
    #                        [-t_ac[1],   t_ac[0],         0]])
    #     s_x = -t_cdc[0]
    #     s_y = -t_cdc[1]
    #     s_z = -t_cdc[2]
    #     q = RR.from_matrix(R_cdc).as_quat()
    #     # q = RR.from_matrix(R_ccd).as_quat()
    #     rotvec = RR.from_quat(q).as_rotvec()
    #     axis, angle = self.quat_to_axis_angle(q)
    #     axis_s = np.array([[       0, -axis[2],  axis[1]],
    #                        [ axis[2],        0, -axis[0]],
    #                        [-axis[1],  axis[0],        0]])
    #     L_theta_u = np.eye(3) - angle/2 * axis_s + (1 - np.sinc(angle)/(np.sinc(angle/2)*np.sinc(angle/2))) * np.dot(axis_s,axis_s)

    #     Le_inv = np.vstack((     np.hstack((-np.eye(3),  np.dot(t_ac_s,np.linalg.inv(L_theta_u)))),
    #                         np.hstack((np.zeros((3,3)),                np.linalg.inv(L_theta_u)))))
    #     error = np.array([s_x, s_y, s_z, rotvec[0], rotvec[1], rotvec[2]])
    #     # error = np.array([s_x, s_y, s_z, 0.0,0.0,0.0])
    #     # print(error)

    #     u = self.kp * error

    #     R_bc = self.getBodyMat("0_end_effector")@self.R_ec
    #     R_cb = np.linalg.inv(R_bc)

    #     if np.linalg.norm(error) < gain :
    #         return np.zeros(6)
    #     else:
    #         print("has error")
    #         v = - np.dot(Le_inv , u)
    #         Rv = np.vstack((np.hstack((R_cb, np.zeros((3,3)))),
    #                         np.hstack((np.zeros((3,3)), R_cb))))
    #         v = np.dot(Rv,v)
    #         # v[3] = 0
    #         # v[4] = 0
    #         # v[5] = 0

    #         v[0] = 0
    #         v[1] = 0
    #         v[2] = 0

    #         return v
        
    def reset(self):
        self.is_detect = True
        self.detect_buf = []
        super().reset()


    def step(self,actions):
        # self.rvec,self.tvec = self.aruco_detection() 
        self.rv,self.tv, detect_res = self.mj_aruco_detection() #from cam frame, coordinate refered to aruco
        self.detect_buf.append(detect_res)
        if len(self.detect_buf) == 300:
            if not any(self.detect_buf) :
                self.is_detect = False
            else:
                self.is_detect = True
            self.detect_buf = []



        # print("in step self.tv=",self.tv)
        # print("in step self.rv=",self.rv)
        if self.Detected_already:
            vel = self.pbvs_controller_origin(1E-04, actions, self.rv,self.tv)
            # vel = actions from PPO
            # vel = self.pbvs_controller_PD(1E-04, 0.9, 0.01,self.rv,self.tv)
            action_cam = self.Vel2qpos(vel,self.arm)
            
        else:
            action_cam = np.array([0.0,0.0,0.0,1.5708,0.0,0.0,0.0])
        action_aruco = np.array([0.0,0.0,0.0,1.5708,0.0,0.0,0.0])
        action = np.concatenate(([0.0], action_cam, action_aruco), axis=0)
       
        print("if detected==", self.Detected_already)
        # action = np.array([0.0,
        #                    0.0,0.0,0.0,1.67,0.0,0.0,0.0,
        #                    0.0,0.0,0.0,1.57,0.0,0.05,0.0])
        
        super().step(action)
      
    
if __name__ == "__main__":
    try:
        from roboimi.assets.robots.bidiana_med_vs import BiDianaMedVS
        env = DualVisualServoEnv(
            robot=BiDianaMedVS(),
            is_render=True,
            control_freq=200,
            is_interpolate=True,
        )
        env.reset()
        env.cam_start()
        
        for t in range(int(1e6)):
            # action = np.array([0.63116, -0.09768533, 0.46947228, 1, 1, 0, 0.63116, -0.09768533, 0.46947228, 1, 0, 0])
            # action = np.zeros(12)
            # action[0] = 1.0
            # action[6] = 1.0
            # env.step(action)
            env.step()
            if env.is_render:
                env.render()
    except:
        KeyboardInterrupt