#!/usr/bin/env python3
import rclpy
from rclpy.node import Node
from sensor_msgs.msg import PointCloud2
from geometry_msgs.msg import PointStamped
from std_msgs.msg import Float32
from tf2_ros import Buffer, TransformListener
import tf2_geometry_msgs
import numpy as np
import math
import sensor_msgs_py.point_cloud2 as pc2

class Cluster5Tracker(Node):
    def __init__(self):
        super().__init__('cluster_tracker')

        self.source_frame = 'lidar_base'
        self.target_frame = 'livox_base'

        self.buf = []

        # Subscribe to point cloud of cluster 5
        self.sub = self.create_subscription(
            PointCloud2,
            '/tracking/target',
            self.cb,
            10
        )

        # Publishers
        self.centroid_pub = self.create_publisher(PointStamped, '/tracking/target_pos', 10)
        self.angle_pub = self.create_publisher(Float32, '/nucleo/stepper0/target', 10)

        # TF
        self.tf_buffer = Buffer()
        self.tf_listener = TransformListener(self.tf_buffer, self)

    def cb(self, msg: PointCloud2):
        # Extract raw x, y, z as a regular list
        points = [[p[0], p[1], p[2]] for p in pc2.read_points(msg, field_names=("x","y","z"), skip_nans=True)]
        if len(points) == 0:
            self.get_logger().warn("No points in cluster5 PointCloud2")
            return

        arr = np.array(points, dtype=np.float32)
        centroid = np.mean(arr, axis=0)

        pt = PointStamped()
        pt.header.frame_id = msg.header.frame_id
        pt.header.stamp = msg.header.stamp
        pt.point.x, pt.point.y, pt.point.z = centroid

        try:
            # Transform to livox_base
            t = self.tf_buffer.lookup_transform(
                self.target_frame,
                pt.header.frame_id,
                pt.header.stamp,
                timeout=rclpy.duration.Duration(seconds=1.0)
            )
            pt_transformed = tf2_geometry_msgs.do_transform_point(pt, t)
        except Exception as e:
            self.get_logger().warn(f"TF lookup failed: {e}")
            return

        # Publish centroid
        self.centroid_pub.publish(pt_transformed)

        # Compute yaw angle in degrees
        x, y = pt_transformed.point.x, pt_transformed.point.y
        yaw_deg = -math.degrees(math.atan2(y, x))

        self.buf.append(yaw_deg)
        if len(self.buf) > 15:
            self.buf.pop(0)

        angle_msg = Float32()
        angle_msg.data = sum(self.buf) / len(self.buf)
        self.angle_pub.publish(angle_msg)

        #self.get_logger().info(f"Centroid: ({x:.2f}, {y:.2f}, {pt_transformed.point.z:.2f}), Yaw: {yaw_deg:.2f}°")

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

if __name__ == "__main__":
    main()
