#!/usr/bin/env python3
import rclpy
from rclpy.node import Node
from std_msgs.msg import Float32
import matplotlib.pyplot as plt
import matplotlib.animation as animation
import threading

class MultiPlotter(Node):
    def __init__(self):
        super().__init__("multi_plotter")

        # Data buffers
        self.angle = []
        self.velocity = []
        self.acceleration = []

        # Subscriptions
        self.create_subscription(Float32, '/nucleo/stepper0/position',
                                 self.angle_callback, 10)
        self.create_subscription(Float32, '/nucleo/stepper0/velocity',
                                 self.velocity_callback, 10)
        self.create_subscription(Float32, '/nucleo/stepper0/acceleration',
                                 self.accel_callback, 10)

    def angle_callback(self, msg):
        self.angle.append(msg.data)
        self._trim()

    def velocity_callback(self, msg):
        self.velocity.append(msg.data)
        self._trim()

    def accel_callback(self, msg):
        self.acceleration.append(msg.data)
        self._trim()

    def _trim(self, limit=2000):
        """Keep buffers short so memory stays small."""
        if len(self.angle) > limit:
            self.angle.pop(0)
        if len(self.velocity) > limit:
            self.velocity.pop(0)
        if len(self.acceleration) > limit:
            self.acceleration.pop(0)


def animate(i, node, axes):
    # Clear each subplot
    for ax in axes:
        ax.cla()

    # Angle
    if len(node.angle) > 0:
        axes[0].plot(node.angle, color='red')
    axes[0].set_title("Angle")
    axes[0].grid(True)

    # Velocity
    if len(node.velocity) > 0:
        axes[1].plot(node.velocity, color='green')
    axes[1].set_title("Velocity")
    axes[1].grid(True)

    # Acceleration
    if len(node.acceleration) > 0:
        axes[2].plot(node.acceleration, color='blue')
    axes[2].set_title("Acceleration")
    axes[2].grid(True)

    # Shared axis labels
    axes[2].set_xlabel("Samples")


def main(args=None):
    rclpy.init(args=args)
    node = MultiPlotter()

    # Start ROS spinning in the background
    thread = threading.Thread(target=rclpy.spin, args=(node,), daemon=True)
    thread.start()

    # Create three subplots
    fig, axes = plt.subplots(3, 1, figsize=(8, 8), sharex=True)

    # Persist animation object
    global ani
    ani = animation.FuncAnimation(
        fig,
        animate,
        fargs=(node, axes),
        interval=16
    )

    plt.tight_layout()
    plt.show()

    node.destroy_node()
    rclpy.shutdown()


if __name__ == "__main__":
    main()