import numpy as np
from roboimi.assets.robots.bidiana_med_vs import BiDianaMedVS
from roboimi.envs.visual_servo.double_vs_base import DualVisualServoEnv
import roboimi.utils.KDL_utils.transform as T
import cv2
from scipy.spatial.transform import Rotation as RR
import rclpy
from rclpy.node import Node
from geometry_msgs.msg import PoseArray, Pose
from tf2_ros import TransformBroadcaster,StaticTransformBroadcaster
from geometry_msgs.msg import TransformStamped


class DualVisualServoRosEnv(DualVisualServoEnv):
    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
        )
        rclpy.init()
        self.node = Node('visual_servo_tf_node')
        self.tf_broadcaster_aruco_detect = TransformBroadcaster(self.node)
        self.tf_broadcaster_ee = TransformBroadcaster(self.node)
        self.tf_broadcaster_aruco_ee = TransformBroadcaster(self.node)
        self.tf_broadcaster_mj_cam = TransformBroadcaster(self.node)

        self.tf_broadcaster_base = StaticTransformBroadcaster(self.node)
        self.tf_broadcaster_aruco_base = StaticTransformBroadcaster(self.node)
        self.tf_broadcaster_cam = StaticTransformBroadcaster(self.node)
        self.tf_broadcaster_cam_desire = StaticTransformBroadcaster(self.node)
        
        
    def step(self):
        super().step()
        
        #-----------------------------------------------------------
        t_base = TransformStamped()
        t_base.header.stamp = self.node.get_clock().now().to_msg()
        t_base.header.frame_id = 'world'
        t_base.child_frame_id = 'base_link'
        pos_bs = self.getBodyPos("base_link")
        quat_bs = self.getBodyQuat("base_link")

        t_base.transform.translation.x = pos_bs[0]
        t_base.transform.translation.y = pos_bs[1]
        t_base.transform.translation.z = pos_bs[2]

        t_base.transform.rotation.x = quat_bs[1]
        t_base.transform.rotation.y = quat_bs[2]
        t_base.transform.rotation.z = quat_bs[3]
        t_base.transform.rotation.w = quat_bs[0]
        self.tf_broadcaster_base.sendTransform(t_base)

        #-------------------------------------------------------------
        t_base_a = TransformStamped()
        t_base_a.header.stamp = self.node.get_clock().now().to_msg()
        t_base_a.header.frame_id = 'world'
        t_base_a.child_frame_id = 'aruco_base_link'
        pos_bs_a = self.getBodyPos("base_link_aruco")
        quat_bs_a = self.getBodyQuat("base_link_aruco")

        t_base_a.transform.translation.x = pos_bs_a[0]
        t_base_a.transform.translation.y = pos_bs_a[1]
        t_base_a.transform.translation.z = pos_bs_a[2]

        t_base_a.transform.rotation.x = quat_bs_a[1]
        t_base_a.transform.rotation.y = quat_bs_a[2]
        t_base_a.transform.rotation.z = quat_bs_a[3]
        t_base_a.transform.rotation.w = quat_bs_a[0]
        self.tf_broadcaster_aruco_base.sendTransform(t_base_a)

        #--------------------------------------------------------------
        t_aruco_ee = TransformStamped()
        t_aruco_ee.header.stamp = self.node.get_clock().now().to_msg()
        t_aruco_ee.header.frame_id = 'world'
        t_aruco_ee.child_frame_id = 'aruco_ee_link'
        pos_aruco_ee = self.getBodyPos("aruco")
        quat_aruco_ee = self.getBodyQuat("aruco")

        t_aruco_ee.transform.translation.x = pos_aruco_ee[0]
        t_aruco_ee.transform.translation.y = pos_aruco_ee[1]
        t_aruco_ee.transform.translation.z = pos_aruco_ee[2]
        t_aruco_ee.transform.rotation.x = quat_aruco_ee[1]
        t_aruco_ee.transform.rotation.y = quat_aruco_ee[2]
        t_aruco_ee.transform.rotation.z = quat_aruco_ee[3]
        t_aruco_ee.transform.rotation.w = quat_aruco_ee[0]
        self.tf_broadcaster_aruco_ee.sendTransform(t_aruco_ee)

        #-------------------------------------------------------------
        t_robot = TransformStamped()
        t_robot.header.stamp = self.node.get_clock().now().to_msg()
        t_robot.header.frame_id = 'world'
        t_robot.child_frame_id = 'ee_link'
        pos_ee = self.getBodyPos("0_end_effector")
        quat_ee = self.getBodyQuat("0_end_effector")

        t_robot.transform.translation.x = pos_ee[0]
        t_robot.transform.translation.y = pos_ee[1]
        t_robot.transform.translation.z = pos_ee[2]

        t_robot.transform.rotation.x = quat_ee[1]
        t_robot.transform.rotation.y = quat_ee[2]
        t_robot.transform.rotation.z = quat_ee[3]
        t_robot.transform.rotation.w = quat_ee[0]
        self.tf_broadcaster_ee.sendTransform(t_robot)

        t_mj_cam = TransformStamped()
        t_mj_cam.header.stamp = self.node.get_clock().now().to_msg()
        t_mj_cam.header.frame_id = 'base_link'
        t_mj_cam.child_frame_id = 'mj_cam_link'
        pos_mj_cam = self.getBodyPos("0_cam")
        quat_mj_cam = self.getBodyQuat("0_cam")

        t_mj_cam.transform.translation.x = pos_mj_cam[0]
        t_mj_cam.transform.translation.y = pos_mj_cam[1]
        t_mj_cam.transform.translation.z = pos_mj_cam[2]
        t_mj_cam.transform.rotation.x = quat_mj_cam[1]
        t_mj_cam.transform.rotation.y = quat_mj_cam[2]
        t_mj_cam.transform.rotation.z = quat_mj_cam[3]
        t_mj_cam.transform.rotation.w = quat_mj_cam[0]
        self.tf_broadcaster_mj_cam.sendTransform(t_mj_cam)

        # #---------------------------------------------------
        t_cam = TransformStamped()
        t_cam.header.stamp = self.node.get_clock().now().to_msg()
        t_cam.header.frame_id = 'ee_link'
        t_cam.child_frame_id = 'cam_link'
        t_cam.transform.translation.x = self.t_ec[0]
        t_cam.transform.translation.y = self.t_ec[1]
        t_cam.transform.translation.z = self.t_ec[2]

        t_cam.transform.rotation.x = self.quat_ec[1]
        t_cam.transform.rotation.y = self.quat_ec[2]
        t_cam.transform.rotation.z = self.quat_ec[3]
        t_cam.transform.rotation.w = self.quat_ec[0]
        self.tf_broadcaster_cam.sendTransform(t_cam)
        
        # #---------------------------------------------------
        if self.Detected_already:
            t_aruco = TransformStamped()
            t_aruco.header.stamp = self.node.get_clock().now().to_msg()
            t_aruco.header.frame_id = 'cam_link'
            t_aruco.child_frame_id = 'aruco_link'
            # rot_matrix = np.eye(3)
            # rot_matrix = cv2.Rodrigues(np.array(self.rvec))[0]
            # quat = RR.from_matrix(rot_matrix).as_quat()
            quat = RR.from_rotvec(self.rv).as_quat()
            
            t_aruco.transform.translation.x = self.tv[0]
            t_aruco.transform.translation.y = self.tv[1]
            t_aruco.transform.translation.z = self.tv[2]
            # print("self.tv=",self.tv)
            # print("self.rv=",self.rv)
            t_aruco.transform.rotation.x = quat[0]
            t_aruco.transform.rotation.y = quat[1]
            t_aruco.transform.rotation.z = quat[2]
            t_aruco.transform.rotation.w = quat[3]
            self.tf_broadcaster_aruco_detect.sendTransform(t_aruco)

        # #----------------------------------------------------
            t_cam_des = TransformStamped()
            t_cam_des.header.stamp = self.node.get_clock().now().to_msg()
            t_cam_des.header.frame_id = 'aruco_link'
            t_cam_des.child_frame_id = 'cam_des_link'
            # desire_p = self.desire_p.reshape(1,3)
            desire_p = self.T_acd[:3,3]
            desire_mat = self.T_acd[:3,:3]
            # print("T_acd =",self.T_acd)
            desire_quat = RR.from_matrix(desire_mat).as_quat()
            t_cam_des.transform.translation.x = desire_p[0]
            t_cam_des.transform.translation.y = desire_p[1]
            t_cam_des.transform.translation.z = desire_p[2]

            t_cam_des.transform.rotation.x = desire_quat[0]
            t_cam_des.transform.rotation.y = desire_quat[1]
            t_cam_des.transform.rotation.z = desire_quat[2]
            t_cam_des.transform.rotation.w = desire_quat[3]
            self.tf_broadcaster_cam_desire.sendTransform(t_cam_des)

        
if __name__ == "__main__":
    try:
        
        env = DualVisualServoRosEnv(
            robot=BiDianaMedVS(),
            is_render=True,
            control_freq=200,
            is_interpolate=True,
        )
       
        env.reset()
        env.cam_start()
        for t in range(int(1e6)):
            env.step()
            if env.is_render:
                env.render()
    except:
        KeyboardInterrupt