#!/usr/bin/env python3
import rclpy
from rclpy.node import Node
from std_msgs.msg import Float32
from geometry_msgs.msg import TransformStamped
import tf2_ros
import math

def euler_to_quaternion(roll, pitch, yaw):
    """Convert Euler angles to quaternion (x, y, z, w)."""
    cr = math.cos(roll / 2)
    sr = math.sin(roll / 2)
    cp = math.cos(pitch / 2)
    sp = math.sin(pitch / 2)
    cy = math.cos(yaw / 2)
    sy = math.sin(yaw / 2)

    w = cr * cp * cy + sr * sp * sy
    x = sr * cp * cy - cr * sp * sy
    y = cr * sp * cy + sr * cp * sy
    z = cr * cp * sy - sr * sp * cy
    return x, y, z, w


class ServoTfNode(Node):
    """Publishes TF for a servo based on its current position."""

    def __init__(self):
        super().__init__("stepper_tf_node")

        # Subscribe to current servo position
        self.pos_sub = self.create_subscription(
            Float32,
            "/nucleo/stepper0/position",
            self.pos_callback,
            10
        )

        # TF broadcaster
        self.br = tf2_ros.TransformBroadcaster(self)
        self.get_logger().info("Servo TF node started, subscribing to /nucleo/stepper0/position")

    def pos_callback(self, msg: Float32):
        """Publish a TF corresponding to the servo position."""
        angle_deg = msg.data  # get float from subscriber

        t = TransformStamped()
        #t.header.stamp = (self.get_clock().now() - rclpy.duration.Duration(seconds=0.1)).to_msg()
        t.header.stamp = self.get_clock().now().to_msg()
        t.header.frame_id = "livox_base"
        t.child_frame_id = "frame_default"

        t.transform.translation.x = 0.0
        t.transform.translation.y = 0.0
        t.transform.translation.z = 0.0

        # Convert angle to yaw (radians) and then to quaternion
        yaw = angle_deg * math.pi / 180.0
        qx, qy, qz, qw = euler_to_quaternion(0, 0, yaw)

        t.transform.rotation.x = qx
        t.transform.rotation.y = qy
        t.transform.rotation.z = qz
        t.transform.rotation.w = qw

        # Publish the transform
        self.br.sendTransform(t)


def main(args=None):
    rclpy.init(args=args)
    node = ServoTfNode()
    try:
        rclpy.spin(node)
    except KeyboardInterrupt:
        pass
    finally:
        node.destroy_node()
        rclpy.shutdown()


if __name__ == "__main__":
    main()

