"""Tests for IMU calibration node."""

import pytest
import numpy as np
from sensor_msgs.msg import Imu
from geometry_msgs.msg import Quaternion
import rclpy
from rclpy.time import Time

from saltybot_imu_calibration.imu_calibration_node import IMUCalibrationNode


@pytest.fixture
def rclpy_fixture():
    rclpy.init()
    yield
    rclpy.shutdown()


@pytest.fixture
def node(rclpy_fixture):
    node = IMUCalibrationNode()
    yield node
    node.destroy_node()


class TestInit:
    def test_node_initialization(self, node):
        assert node.calibration_samples == 100
        assert node.is_calibrated is False
        assert node.calibrating is False


class TestCalibration:
    def test_calibration_samples(self, node):
        node.calibration_samples = 3
        node.gyro_samples.maxlen = 3
        node.accel_samples.maxlen = 3
        node.calibrating = True

        for i in range(3):
            node.gyro_samples.append(np.array([0.1, 0.2, 0.3]))
            node.accel_samples.append(np.array([0.0, 0.0, 9.81]))

        node._compute_calibration()

        assert node.is_calibrated is True
        assert len(node.gyro_samples) == 0


class TestCorrection:
    def test_imu_correction(self, node):
        node.gyro_bias = np.array([0.1, 0.2, 0.3])
        node.accel_bias = np.array([0.0, 0.0, 0.1])

        msg = Imu()
        msg.header.stamp = Time().to_msg()
        msg.header.frame_id = "imu_link"
        msg.angular_velocity.x = 0.11
        msg.angular_velocity.y = 0.22
        msg.angular_velocity.z = 0.33
        msg.linear_acceleration.x = 0.0
        msg.linear_acceleration.y = 0.0
        msg.linear_acceleration.z = 9.91
        msg.orientation = Quaternion(x=0, y=0, z=0, w=1)

        node._publish_calibrated(msg)
