import rclpy
from rclpy.node import Node
from sensor_msgs.msg import JointState
from trajectory_msgs.msg import JointTrajectory

import numpy as np
import time

class RobotNode(Node):

    def __init__(self):
        super().__init__('robot_node')

        self.l1 = 0.5
        self.l2 = 0.3
        self.x = self.l1 + self.l2
        self.y = 0.0
        self.theta1 = 0.0
        self.theta2 = 0.0

        self.publisher = self.create_publisher(
            JointState, 
            '/joint_states', 
            1
        )

        timer_period = 0.01  # seconds
        self.timer = self.create_timer(timer_period, self.timer_callback)
        self.i = 0

        self.subcriber = self.create_subscription(
            JointTrajectory,
            '/scaled_joint_trajectory_controller/joint_trajectory',
            self.joint_trajectory_callback,
            1
        )

    def timer_callback(self):
        msg = JointState()
        msg.name = ['joint1', 'joint2']
        msg.header.stamp = self.get_clock().now().to_msg()
        msg.position = [self.theta1, self.theta2]
        msg.velocity = [self.x, self.y]
        self.publisher.publish(msg)

    def joint_trajectory_callback(self, msg):
        joint_names = msg.joint_names
        prev_timetag = 0
        for point in msg.points:
            duration = point.time_from_start.sec + point.time_from_start.nanosec / 1e9 - prev_timetag
            prev_timetag = point.time_from_start.sec + point.time_from_start.nanosec / 1e9
            
            steps = int(duration * 10)  # Update in 10 steps per second

            for i in range(steps):
                self.theta1 += (point.positions[0] - self.theta1) / steps
                self.theta2 += (point.positions[1] - self.theta2) / steps
                self.x = self.l1 * np.cos(self.theta1) + self.l2 * np.cos(self.theta1 + self.theta2)
                self.y = self.l1 * np.sin(self.theta1) + self.l2 * np.sin(self.theta1 + self.theta2)
                msg = JointState()
                msg.name = ['joint1', 'joint2']
                msg.header.stamp = self.get_clock().now().to_msg()
                msg.position = [self.theta1, self.theta2]
                msg.velocity = [self.x, self.y]
                self.publisher.publish(msg)
                time.sleep(0.1)


def main(args=None):
    rclpy.init(args=args)
    node = RobotNode()
    rclpy.spin(node)
    rclpy.shutdown()