#!/usr/bin/python3

import rostest
import rospy
import roslaunch

import threading
from math import sqrt
import socket
import unittest
from sbp.imu import MsgImuRaw
from sbp.mag import MsgMagRaw
from sbp.navigation import MsgPosLLH

from sensor_msgs.msg import NavSatFix, NavSatStatus, Imu, MagneticField

NODE = 'piksi_multi_rtk'
ACC_SENS = 0.0001
GYRO_SENS = 0.0002
MSGS_TO_TEST = 100


class SingleBufferSubscriber:
    def __init__(self, topic, msg_type):
        def cb(msg):
            with self.lock:
                self.msg = msg

        self.lock = threading.Lock()
        self.msg = None
        self.sub = rospy.Subscriber(topic, msg_type, callback=cb, queue_size=1)

    def ready(self):
        return self.sub.get_num_connections() > 0

    def get_msg(self, timeout=10):
        final = rospy.Time.now() + rospy.Duration(timeout)
        while rospy.Time.now() <= final:
            with self.lock:
                if self.msg is not None:
                    msg = self.msg
                    self.msg = None
                    return msg
            rospy.sleep(0.01)
        raise TimeoutError('No message detected')


class TestPiksiMultiRTK(unittest.TestCase):

    @classmethod
    def setUpClass(cls):
        rospy.init_node(f'test_{NODE}')

        rospy.set_param(f'{NODE}/tcp_addr', '127.0.0.1')
        rospy.set_param(f'{NODE}/imu_frame', 'imu')
        rospy.set_param(f'{NODE}/mag_frame', 'mag')
        rospy.set_param(f'{NODE}/imu_acc_sens', ACC_SENS)
        rospy.set_param(f'{NODE}/imu_gyro_sens', GYRO_SENS)

        node = roslaunch.core.Node('piksi_multi_rtk', NODE, name=NODE)

        launch = roslaunch.scriptapi.ROSLaunch()
        launch.start()

        cls.process = launch.launch(node)

        cls.sub_imu = SingleBufferSubscriber('imu/raw', Imu)

        cls.server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)

        # Allow reuse of old sockets still in TIME_WAIT state
        cls.server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)

        # Bind the socket to the given port, an empty string as first element
        # in the tuple signals that this should be a server socket (for a
        # client socket this element would have specified an ip to connect to)
        cls.server.bind(('', 55555))

        # Make the socket ready to accept connections. Five clients can be in
        # the queue before connections get rejected
        cls.server.settimeout(10)
        cls.server.listen(1)

        cls.client = cls.server.accept()[0]
        cls.client.settimeout(10)

        cls.sub_mag = SingleBufferSubscriber('mag/raw', MagneticField)
        cls.sub_fix = SingleBufferSubscriber('gnss/fix', NavSatFix)

        def topics_are_published():
            return (cls.sub_imu.ready()
                    and cls.sub_imu.ready()
                    and cls.sub_imu.ready())

        final = rospy.Time.now() + rospy.Duration(10)
        while not topics_are_published() and rospy.Time.now() <= final:
            rospy.sleep(0.1)
        cls.topics_are_published = topics_are_published()

        # Just some extra delay in case the system needs it idk...
        rospy.sleep(1.0)

    def setUp(self):
        if not TestPiksiMultiRTK.topics_are_published:
            self.fail('Topics were not published in time')

    def _navsatfix_to_sbp(self, fix):
        return MsgPosLLH(
            sender=1,
            tow=0,
            lat=fix.latitude,
            lon=fix.longitude,
            height=fix.altitude,
            h_accuracy=int(round(sqrt(fix.position_covariance[0]) * 1e3)),
            v_accuracy=int(round(sqrt(fix.position_covariance[8]) * 1e3)),
            n_sats=1,
            flags=0
        )

    def _imu_to_sbp(self, imu):
        return MsgImuRaw(
            tow=0,
            tow_f=0,
            acc_x=int(round(imu.linear_acceleration.x / ACC_SENS)),
            acc_y=int(round(imu.linear_acceleration.y / ACC_SENS)),
            acc_z=int(round(imu.linear_acceleration.z / ACC_SENS)),
            gyr_x=int(round(imu.angular_velocity.x / GYRO_SENS)),
            gyr_y=int(round(imu.angular_velocity.y / GYRO_SENS)),
            gyr_z=int(round(imu.angular_velocity.z / GYRO_SENS))
        )

    def _mag_to_sbp(self, mag):
        return MsgMagRaw(
            tow=0,
            tow_f=0,
            mag_x=int(round(mag.magnetic_field.x * 1e6)),
            mag_y=int(round(mag.magnetic_field.y * 1e6)),
            mag_z=int(round(mag.magnetic_field.z * 1e6))
        )

    def test_llh_msgs_are_published(self):
        fix = NavSatFix()
        fix.position_covariance = [4, 0, 0,
                                   0, 4, 0,
                                   0, 0, 5]

        for i in range(MSGS_TO_TEST):
            f = (2.0 * i - MSGS_TO_TEST) / MSGS_TO_TEST
            fix.altitude = f * 100
            fix.latitude = f * 80
            fix.longitude = (1 - f) * 180

            sbp = self._navsatfix_to_sbp(fix)
            TestPiksiMultiRTK.client.sendall(sbp.to_binary())
            response = TestPiksiMultiRTK.sub_fix.get_msg()

            self.assertAlmostEqual(fix.latitude, response.latitude)
            self.assertAlmostEqual(fix.longitude, response.longitude)
            self.assertAlmostEqual(fix.altitude, response.altitude)
            for i in range(9):
                self.assertAlmostEqual(fix.position_covariance[i],
                                       response.position_covariance[i],
                                       places=3)

    def test_imu_msgs_are_published(self):
        imu = Imu()

        for i in range(MSGS_TO_TEST):
            f = (2.0 * i - MSGS_TO_TEST) / MSGS_TO_TEST
            imu.linear_acceleration.x = f
            imu.linear_acceleration.y = 1 - f
            imu.linear_acceleration.z = 2 * f
            imu.angular_velocity.x = 3 * f
            imu.angular_velocity.y = 2 - f
            imu.angular_velocity.z = -f

            sbp = self._imu_to_sbp(imu)
            TestPiksiMultiRTK.client.sendall(sbp.to_binary())
            response = TestPiksiMultiRTK.sub_imu.get_msg()

            self.assertAlmostEqual(imu.linear_acceleration.x,
                                   response.linear_acceleration.x)
            self.assertAlmostEqual(imu.linear_acceleration.y,
                                   response.linear_acceleration.y)
            self.assertAlmostEqual(imu.linear_acceleration.z,
                                   response.linear_acceleration.z)
            self.assertAlmostEqual(imu.angular_velocity.x,
                                   response.angular_velocity.x)
            self.assertAlmostEqual(imu.angular_velocity.y,
                                   response.angular_velocity.y)
            self.assertAlmostEqual(imu.angular_velocity.z,
                                   response.angular_velocity.z)

    def test_mag_msgs_are_published(self):
        mag = MagneticField()

        for i in range(MSGS_TO_TEST):
            f = (2.0 * i - MSGS_TO_TEST) / MSGS_TO_TEST

            mag.magnetic_field.x = f * 1e-4
            mag.magnetic_field.y = -f * 1e-4
            mag.magnetic_field.z = f * 1e-3

            sbp = self._mag_to_sbp(mag)
            TestPiksiMultiRTK.client.sendall(sbp.to_binary())
            response = TestPiksiMultiRTK.sub_mag.get_msg()

            self.assertAlmostEqual(mag.magnetic_field.x,
                                   response.magnetic_field.x)
            self.assertAlmostEqual(mag.magnetic_field.y,
                                   response.magnetic_field.y)
            self.assertAlmostEqual(mag.magnetic_field.z,
                                   response.magnetic_field.z)

    @classmethod
    def tearDownClass(cls):
        cls.process.stop()

        cls.client.close()
        cls.server.close()


if __name__ == '__main__':
    rostest.rosrun('piksi_multi_rtk', f'test_{NODE}', TestPiksiMultiRTK)
