from math import sin, cos, pi, copysign
import rclpy, random, time, threading
from rclpy.node import Node
from rclpy.qos import QoSProfile
from geometry_msgs.msg import Quaternion
from sensor_msgs.msg import JointState
from tf2_ros import TransformBroadcaster, TransformStamped
import socket
import struct

class StatePublisher(Node):

    def __init__(self):
        super().__init__('state_publisher')

        self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)

        self.sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        self.sock.bind(('0.0.0.0', 9999))
        self.sock.listen(1)

        qos_profile = QoSProfile(depth=10)
        self.joint_pub = self.create_publisher(JointState, 'joint_states', qos_profile)
        self.broadcaster = TransformBroadcaster(self, qos=qos_profile)
        self.nodeName = self.get_name()
        self.get_logger().info("{0} started".format(self.nodeName))

        self.cur_joints = [0.] * 6
        self.cmd_joints = [0.] * 6
        self.rate = 10. / 180. * pi
        self.accel = 10. / 180. * pi
        self.prev_time = time.time()

        threading.Thread(target=self.server_handle).start()

    def server_handle(self):
        while rclpy.ok():
            cli, addr = self.sock.accept()
            cli.settimeout(0)
            cli.setblocking(0)
            print(cli, addr)
            self.prev_time = time.time()
            while rclpy.ok():
                try:
                    msg = cli.recv(2048)
                    if msg == b'':
                        cli.close()
                        print('bailing')
                        break

                    while len(msg) >= 13:
                        if msg[12] == 10:
                            self.cmd_joints = struct.unpack('hhhhhh', msg[:12])
                            self.cmd_joints = [a / 180. * pi for a in self.cmd_joints]
                            msg = msg[13:]
                        print(repr(msg))
                except BlockingIOError as e:
                    pass
                except Exception as e:
                    print('what the fuck')
                    print(e)
                    raise e

                cur_time = time.time()
                dt = cur_time - self.prev_time
                self.prev_time = cur_time
                # print(self.cmd_joints, self.cur_joints)
                # print(dt)

                for i in range(6):
                    self.cur_joints[i] += copysign(self.rate * dt + 0.5 * self.accel * dt * dt, self.cmd_joints[i] - self.cur_joints[i])

                # message declarations
                odom_trans = TransformStamped()
                odom_trans.header.frame_id = 'odom'
                odom_trans.child_frame_id = 'base_link'
                joint_state = JointState()

                # update joint_state
                now = self.get_clock().now()
                joint_state.header.stamp = now.to_msg()
                joint_state.name = ['joint_1', 'joint_2', 'joint_3', 'joint_4', 'joint_5', 'joint_6']
                joint_state.position = self.cur_joints.copy()

                # update transform
                # (moving in a circle with radius=2)
                odom_trans.header.stamp = now.to_msg()
                odom_trans.transform.translation.x = 0.0# cos(angle) * 0.4
                odom_trans.transform.translation.y = 0.0# sin(angle) * 0.4
                odom_trans.transform.translation.z = 0.0
                odom_trans.transform.rotation = \
                    euler_to_quaternion(0, 0, 0) # roll,pitch,yaw

                # send the joint state and transform
                self.joint_pub.publish(joint_state)
                self.broadcaster.sendTransform(odom_trans)

def euler_to_quaternion(roll, pitch, yaw):
    qx = sin(roll/2) * cos(pitch/2) * cos(yaw/2) - cos(roll/2) * sin(pitch/2) * sin(yaw/2)
    qy = cos(roll/2) * sin(pitch/2) * cos(yaw/2) + sin(roll/2) * cos(pitch/2) * sin(yaw/2)
    qz = cos(roll/2) * cos(pitch/2) * sin(yaw/2) - sin(roll/2) * sin(pitch/2) * cos(yaw/2)
    qw = cos(roll/2) * cos(pitch/2) * cos(yaw/2) + sin(roll/2) * sin(pitch/2) * sin(yaw/2)
    return Quaternion(x=qx, y=qy, z=qz, w=qw)

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

    node = StatePublisher()

    rclpy.spin(node)

    node.destroy_node()
    rclpy.shutdown()

if __name__ == '__main__':
    main()
