import numpy as np
from roboimi.envs.single_base import SingleArmEnv
import threading,lcm,rclpy,copy
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

class Diana_Teleop(SingleArmEnv):
    def __init__(self,
                robot=None,
                is_render=True,
                renderer="mujoco_viewer",
                control_freq=200,
                is_interpolate=True,
                is_pd=False,
                ):
        super().__init__(
            robot=robot,
            is_render=is_render,
            renderer=renderer,
            control_freq=control_freq,
            is_interpolate=is_interpolate
        )
        self.qpos = np.zeros(7)
        self.node = Node("diana_sim_publisher")
        self.publisher = self.node.create_publisher(Float32MultiArray, 'mj_joint_states', 10)

        self.sub = LCM_sub("diana_joint_state")
        self.lcm_thread=threading.Thread(target = self.sub.lcm_spin)

    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 reset(self):
        super().reset()

    def preStep(self, action):
        if isinstance(action,np.ndarray) and len(action)==7:
            super().preStep(action)
        else:
            raise AttributeError("size should be 7")

    def step(self, action):
        super().step(action)
        for i in range(7):
            self.qpos[i] = copy.deepcopy(self.robot.robot_data.joint(self.robot.single_arm.joint_index[i]).qpos[0])
        self.publish_joint_angles(self.qpos)
    
    def mj_run(self):
        while True:
            # self.action = np.array([0, 0, 0, 1.57, 0, 0, 0])
            self.action = self.sub.diana_arm
            self.step(self.action)
            if self.is_render:
                self.render()

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]))     
                # if self.diana_arm[i]-self.pre_diana_arm[i] > 0.0001:
                #     self.diana_arm[i] = self.pre_diana_arm[i]+0.0001
                # elif self.diana_arm[i]-self.pre_diana_arm[i] < -0.0001:
                #     self.diana_arm[i] = self.pre_diana_arm[i]-0.0001
                # else:
                #     pass
                self.temp[i] = self.diana_arm[i]
            self.pre_diana_arm = self.temp
            # 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)
    from roboimi.assets.robots.diana_med import DianaMed
    env = Diana_Teleop( 
        robot=DianaMed(),
        renderer='viewer',
        is_render=True,
        control_freq=10,
        is_interpolate=True,
        is_pd=False
    )
    env.reset()
    try:
        env.lcm_thread.start()
        env.mj_run()

    except KeyboardInterrupt:
        print("\nbye")
    finally:
        rclpy.try_shutdown()



if __name__ == '__main__':
    main()


