import numpy as np
from roboimi.envs.single_base import SingleArmEnv
import threading,lcm,rclpy,copy,mujoco,os,time
from roboimi.scripts.lcm_interfaces.diana_joint_states_t import diana_joint_states_t
from std_msgs.msg import Float32MultiArray
from rclpy.node import Node
from rclpy.executors import SingleThreadedExecutor

# 遥操位置控制，即运动学控制时，要关闭xml的重力，即 gravity="disable"
class Diana_Teleop():
    def __init__(self, xml_path):
        self.model = mujoco.MjModel.from_xml_path(xml_path)
        self.data  = mujoco.MjData(self.model)
        self.control_freq = 200
        self.sim_freq = 2000
        self.freq = int(np.floor(self.sim_freq/self.control_freq))
        self.target_qpos = np.array([0.0, 0.0, 0.0, 2.17, 0.0, -0.9, 0.0])

        # keyboard flag
        self.render_paused = True
        self.exit_flag = False
        self.is_render = False
        self.renderer = None
        self.time_interval = 0.005
        self.shared_resource_lock = threading.Lock()

        self.sub = LCM_sub("diana_joint_state")
        self.lcm_thread = threading.Thread(target=self.sub.lcm_spin)
        self.robot_jnt = ["0_j1","0_j2","0_j3","0_j4","0_j5","0_j6","0_j7"]
        self.initRender(True, 'viewer')

        self.arm = np.zeros(7)
        self.node = Node("diana_sim_publisher")
        self.publisher = self.node.create_publisher(Float32MultiArray, 'mj_joint_states', 10)
        self.ros2_thread = threading.Thread(target=self.ros2_publisher)

        # self.timer = self.node.create_timer(0.001, self.timer_callback)  # 每1毫秒调用一次回调
        # self.exe = SingleThreadedExecutor()
        # self.exe.add_node(self.node)


    
    def initRender(self,is_render : bool, viewer : str):
        self.is_render = is_render
        if self.is_render:
            if viewer == "viewer":
                def key_callback(keycode):
                    if keycode == 32: #space
                        self.render_paused = not self.render_paused
                    elif keycode == 256: #escape
                        self.exit_flag = not self.exit_flag

                self.viewer = mujoco.viewer.launch_passive(self.model, self.data,
                                                key_callback=key_callback, show_left_ui=False, show_right_ui=False)
                self.viewer.cam.lookat = np.array([0.4, 0, 0.5])
                self.viewer.cam.azimuth -= 0.005
                with self.viewer.lock():
                    self.viewer.opt.flags[mujoco.mjtVisFlag.mjVIS_CONTACTPOINT] = int(self.data.time % 2)
                print("mujoco viewer init !") 
            else:
                raise ValueError('renderer init failed for some reason.')
        
        else:
            print("No asking for render simulation")
    

    def setArmPos(self,q : np.ndarray):
        for i in range(7):
            self.target_qpos[i] = q[i]
  

    def getArmPos(self):
        q = np.zeros(7)
        for i in range(7):
            q[i] = self.data.joint(self.robot_jnt[i]).qpos[0]
        return q

    def publish_joint_angles(self, angles):
        angles = [float(x) for x in angles]
        msg = Float32MultiArray()
        msg.data = angles
        self.publisher.publish(msg)
        # self.node.get_logger().info(f'Publishing joint angles: {msg.data}')
    
    def timer_callback(self): # ros2线程管理
        self.publish_joint_angles(self.arm)


    def ros2_publisher(self): # thread线程管理
        while True:
            self.publish_joint_angles(self.arm)
            time.sleep(0.001)


    def mj_run(self):
        num = 0
        while True:
            for unit_interval in range(self.freq):
                if self.sub.flag:
                    self.setArmPos(self.sub.diana_arm)    
                    for i in range(7):
                        self.data.joint(self.robot_jnt[i]).qpos[0] = self.target_qpos[i]
                    self.sub.flag = False         
                else:
                    for i in range(7):
                        self.data.joint(self.robot_jnt[i]).qpos[0] = self.target_qpos[i]
                mujoco.mj_forward(self.model, self.data)
                mujoco.mj_step(self.model, self.data)
                self.arm  = self.getArmPos()
                # self.publish_joint_angles(self.arm)
            if self.is_render:
                self.shared_resource_lock.acquire()
                if self.viewer is not None and self.render_paused is True:
                    if self.viewer.is_running() and self.exit_flag is False:
                        self.viewer: mujoco.viewer.Handle
                        self.viewer.sync()
                    else:
                        self.viewer.close() 
                self.shared_resource_lock.release()
            else:
                pass

    
class LCM_sub(object):
    def __init__(self,name : str):
        self.lc = lcm.LCM()
        self.sub_ = self.lc.subscribe(name, self.lcm_handler)
        self.msg = None
        self.diana_arm = np.array([0.0, 0.0, 0.0, 2.17, 0.0, -0.9, 0.0])
        self.pre_diana_arm = np.array([0.0, 0.0, 0.0, 2.17, 0.0, -0.9, 0.0])
        # self.diana_arm = np.array([0.0, 0.0, 0.0, 1.57, 0.0, 0.0, 0.0])
        # self.pre_diana_arm = np.array([0.0, 0.0, 0.0, 1.57, 0.0, 0.0, 0.0])
        self.temp = np.zeros(7)
        self.ulim = np.array([ 2.56,  1.57,  2.56, 3.05,  3.05, 0.523,  3.05])
        self.llim = np.array([-2.56, -1.57, -2.56,  0.0, -3.05, -3.05, -3.05])
        self.flag = False
        
    def lcm_handler(self,channel, data):
        self.msg = diana_joint_states_t.decode(data)
        if self.msg.enabled:
            self.flag = True
            for i in range(7):
                self.diana_arm[i]  = np.maximum(self.llim[i],np.minimum(self.msg.jntpos[i],self.ulim[i]))     
                self.temp[i] = self.diana_arm[i]
                # if self.diana_arm[i]-self.pre_diana_arm[i] > 0.003:
                #     self.diana_arm[i] = self.pre_diana_arm[i]+0.003
                # elif self.diana_arm[i]-self.pre_diana_arm[i] < -0.003:
                #     self.diana_arm[i] = self.pre_diana_arm[i]-0.003
                # else:
                #     pass
            # print("connected!!! \n")
        else:
            pass

    def lcm_spin(self):
        try:
            while True:
                self.lc.handle()
        except KeyboardInterrupt:
            self.lc.unsubscribe(self.sub_)
            print("KeyboardInterrupt !")
        
def main(args=None):
    rclpy.init(args=args)
    CURRENT = os.getcwd()
    scence = CURRENT + '/assets/singlediana.xml'
    print(scence)
    env = Diana_Teleop(scence)
    # 收发线程放在控制线程前
    try:
        env.lcm_thread.start()
        # env.exe.spin()
        env.ros2_thread.start()
        env.mj_run() 
    except KeyboardInterrupt:
        print("\nbye")
    finally:
        rclpy.try_shutdown()



if __name__ == '__main__':
    main()
    

