import rclpy
from rclpy.node import Node
from trajectory_msgs.msg import JointTrajectory
import csv
import os
import roboticstoolbox as rtb

class JointTrajectoryLogger(Node):
    def __init__(self):
        super().__init__('joint_trajectory_logger_minimal')

        # Folder to store CSV files
        self.log_dir = '/BA/joint_trajectories'
        os.makedirs(self.log_dir, exist_ok=True)

        self.log_dir_cart = '/BA/joint_trajectories_cart'
        os.makedirs(self.log_dir_cart, exist_ok=True)

        # Subscriber to JointTrajectory messages
        self.subscription = self.create_subscription(
            JointTrajectory,
            '/scaled_joint_trajectory_controller/joint_trajectory',
            self.trajectory_callback,
            10
        )

        self.counter = 0
        self.robot = rtb.ERobot.URDF('/BA/ur10e.urdf')

    def trajectory_callback(self, msg: JointTrajectory):
        self.counter += 1

        # Create a unique filename
        filename = f"trajectory_{self.counter:04d}.csv"
        file_path = os.path.join(self.log_dir, filename)
        # Header: timestamp + joint names
        header = ['timestamp'] + msg.joint_names
        recive_time = msg.header.stamp.sec + msg.header.stamp.nanosec * 1e-9
        with open(file_path, mode='w', newline='') as file:
            writer = csv.writer(file)
            writer.writerow(header)

            for point in msg.points:
                timestamp = recive_time + point.time_from_start.sec + point.time_from_start.nanosec * 1e-9
                row = [timestamp] + list(point.positions)
                writer.writerow(row)

        filename_cart = f"trajectory_{self.counter:04d}_cart.csv"
        file_path_cart = os.path.join(self.log_dir_cart, filename_cart)
        header = ['timestamp', 'x', 'y', 'z', 'roll', 'pitch', 'yaw']
        # Convert joint angles to cartesian coordinates
        joint_angles = point.positions
        
        # Write to CSV
        with open(file_path_cart, mode='w', newline='') as file:
            writer = csv.writer(file)
            writer.writerow(header)

            for point in msg.points:
                timestamp = recive_time + point.time_from_start.sec + point.time_from_start.nanosec * 1e-9
                T = self.robot.fkine(point.positions)
                x = T.t[0]
                y = T.t[1]
                z = T.t[2]
                roll, pitch, yaw = T.rpy()
                row = [timestamp, x, y, z, roll, pitch, yaw]
                writer.writerow(row)

def main():
    rclpy.init()
    node = JointTrajectoryLogger()
    try:
        rclpy.spin(node)
    except KeyboardInterrupt:
        print("")
    finally:
        node.destroy_node()
        rclpy.shutdown()

if __name__ == '__main__':
    main()