import rclpy
from rclpy.node import Node
from sensor_msgs.msg import JointState
import csv
import os
import roboticstoolbox as rtb

class JointStateLogger(Node):
    def __init__(self):
        super().__init__('joint_state_logger_split')

        # Subscribe to /joint_states topic
        self.subscription = self.create_subscription(
            JointState,
            '/joint_states',
            self.joint_states_callback,
            1000
        )
        self.robot = rtb.ERobot.URDF('/BA/ur10e.urdf')
        # Directory to store logs
        self.log_dir = '/BA/joint_states_logs'
        os.makedirs(self.log_dir, exist_ok=True)

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

    def joint_states_callback(self, msg: JointState):
        # Use ROS time as timestamp
        stamp = msg.header.stamp
        timestamp = f"{stamp.sec}.{str(stamp.nanosec).zfill(9)}"

        for i, name in enumerate(msg.name):
            # File per joint
            file_path = os.path.join(self.log_dir, f"{name}.csv")

            # If file doesn't exist, create with header
            if not os.path.exists(file_path):
                with open(file_path, mode='w', newline='') as f:
                    writer = csv.writer(f)
                    writer.writerow(["timestamp", "position", "velocity", "effort"])

            # Get values if they exist
            position = msg.position[i] if i < len(msg.position) else ''
            velocity = msg.velocity[i] if i < len(msg.velocity) else ''
            effort = msg.effort[i] if i < len(msg.effort) else ''

            # Append to joint-specific CSV
            with open(file_path, mode='a', newline='') as f:
                writer = csv.writer(f)
                writer.writerow([timestamp, position, velocity, effort])
        
        file_path_cart = os.path.join(self.log_dir_cart, "cartesian.csv")
        if not os.path.exists(file_path_cart):
            with open(file_path_cart, mode='w', newline='') as f:
                writer = csv.writer(f)
                writer.writerow(["timestamp", "x", "y", "z", "roll", "pitch", "yaw"])
        with open(file_path_cart, mode='a', newline='') as f:
            writer = csv.writer(f)
            [x,y,z] = self.robot.fkine(list([msg.position[-1]])+list(msg.position[:4])).t
            [roll, pitch, yaw] = self.robot.fkine(list([msg.position[-1]])+list(msg.position[:4])).rpy()
            writer.writerow([timestamp, x, y, z, roll, pitch, yaw])
            

        

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

if __name__ == '__main__':
    main()