import rclpy, re
from rclpy.node import Node
from rclpy.qos import QoSProfile

from sensor_msgs.msg import JointState
from std_msgs.msg import Float32

from math import pi

from serial import Serial

J1_COUNTS_PER_RAD = 1 / (2 * pi) * 400 * 10 * 4
J2_COUNTS_PER_RAD = 1 / (2 * pi) * 400 * 50
J3_COUNTS_PER_RAD = 1 / (2 * pi) * 400 * 50
J4_COUNTS_PER_RAD = 1 / (2 * pi) * 400 * (13 + 212 / 289) * 2.8
J5_COUNTS_PER_RAD = 1 / (2 * pi) * 800 * 10
J6_COUNTS_PER_RAD = 1 / (2 * pi) * 400 * (19 + 38 / 187)

class CobotMove(Node):

    def __init__(self):
        super().__init__('cobot_move')
        # Set up a logger for the node
        self._logger = self.get_logger()
        self._logger.info('COBOT init')

        # Connect to the serial port on the teensy
        self._serial = Serial('/dev/ttyACM0', 115200)

        # Set initial joint values to all 0s
        self._cur = [0.0] * 6

        # create a timer to update all the joint angles at once
        self._timer = self.create_timer(1 / 40, self._dosomething)

        # Set joint angles independantly
        base_topic = 'move/'
        self.create_subscription(Float32, base_topic+'j1', self._move_j1_cb, rclpy.qos.qos_profile_sensor_data)
        self.create_subscription(Float32, base_topic+'j2', self._move_j2_cb, rclpy.qos.qos_profile_sensor_data)
        self.create_subscription(Float32, base_topic+'j3', self._move_j3_cb, rclpy.qos.qos_profile_sensor_data)
        self.create_subscription(Float32, base_topic+'j4', self._move_j4_cb, rclpy.qos.qos_profile_sensor_data)
        self.create_subscription(Float32, base_topic+'j5', self._move_j5_cb, rclpy.qos.qos_profile_sensor_data)
        self.create_subscription(Float32, base_topic+'j6', self._move_j6_cb, rclpy.qos.qos_profile_sensor_data)

        qos = QoSProfile(depth=10)

        # Publishing the values so other nodes can read them (so they don't need to keep track of current positions)
        publish_base_topic = 'update/'
        self.s_j1 = self.create_publisher(Float32, publish_base_topic+'j1', qos)
        self.s_j2 = self.create_publisher(Float32, publish_base_topic+'j2', qos)
        self.s_j3 = self.create_publisher(Float32, publish_base_topic+'j3', qos)
        self.s_j4 = self.create_publisher(Float32, publish_base_topic+'j4', qos)
        self.s_j5 = self.create_publisher(Float32, publish_base_topic+'j5', qos)
        self.s_j6 = self.create_publisher(Float32, publish_base_topic+'j6', qos)

        # Set the initial position of the robot to 0
        self._serial.write(b'STA0.0.1\n')
        print(self._serial.readline())

    # Callbacks for move/jX subscriptions
    def _move_j1_cb(self, msg):
        self._cur[0] = int(msg.data * J1_COUNTS_PER_RAD)

    def _move_j2_cb(self, msg):
        self._cur[1] = int(msg.data * J2_COUNTS_PER_RAD)

    def _move_j3_cb(self, msg):
        self._cur[2] = int(msg.data * J3_COUNTS_PER_RAD)

    def _move_j4_cb(self, msg):
        self._cur[3] = int(msg.data * J4_COUNTS_PER_RAD)

    def _move_j5_cb(self, msg):
        self._cur[4] = int(msg.data * J5_COUNTS_PER_RAD)

    def _move_j6_cb(self, msg):
        self._cur[5] = int(msg.data * J6_COUNTS_PER_RAD)

    # Send all current positions to the teensy
    def _dosomething(self):
        self._serial.write(f'MTA{self._cur[0]}B{self._cur[1]}C{self._cur[2]}D{self._cur[3]}E{self._cur[4]}F{self._cur[5]}\n'.encode('UTF-8'))

        # Get a response from line above of positions
        line = self._serial.readline().decode('UTF-8')

        # Parse the line into each angle
        blocks = re.findall(r'[ABCDEF]-*\d+', line)
        a = float(blocks[0][1:]) / J1_COUNTS_PER_RAD
        b = float(blocks[1][1:]) / J2_COUNTS_PER_RAD
        c = float(blocks[2][1:]) / J3_COUNTS_PER_RAD
        d = float(blocks[3][1:]) / J4_COUNTS_PER_RAD
        e = float(blocks[4][1:]) / J5_COUNTS_PER_RAD
        f = float(blocks[5][1:]) / J6_COUNTS_PER_RAD

        a = float(self._cur[0] / J1_COUNTS_PER_RAD)
        b = float(self._cur[1] / J2_COUNTS_PER_RAD)
        c = float(self._cur[2] / J3_COUNTS_PER_RAD)
        d = float(self._cur[3] / J4_COUNTS_PER_RAD)
        e = float(self._cur[4] / J5_COUNTS_PER_RAD)
        f = float(self._cur[5] / J6_COUNTS_PER_RAD)

        # Publish those to the update/jX subscribers
        self.s_j1.publish(Float32(data=a))
        self.s_j2.publish(Float32(data=b))
        self.s_j3.publish(Float32(data=c))
        self.s_j4.publish(Float32(data=d))
        self.s_j5.publish(Float32(data=e))
        self.s_j6.publish(Float32(data=f))


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

    # start the node
    cobot_move = CobotMove()
    rclpy.spin(cobot_move)

    # kill the node
    cobot_move.destroy_node()
    rclpy.shutdown()


if __name__ == '__main__':
    main()

