import unittest

import numpy as np

from scripts import imu_static_calibrator as calibrator


class ImuStaticCalibratorTests(unittest.TestCase):
    def _state(self, enter_seconds=0.01):
        config = calibrator.StaticCorrectionConfig(enter_seconds=enter_seconds)
        return calibrator.initialize(config, initial_yaw_bias_z_dps=0.0)

    def test_stationary_samples_enter_static_and_estimate_z_bias(self):
        state = self._state()

        for _ in range(7):
            is_static = calibrator.step(
                state,
                0.002,
                np.array([0.0, 0.0, 1.1]),
                np.array([0.0, 0.0, 0.12]),
                np.zeros(2),
            )

        self.assertTrue(is_static)
        self.assertAlmostEqual(state.active_yaw_bias_z_dps, 0.12, places=9)

    def test_rotation_above_threshold_never_enters_static(self):
        state = self._state()

        for _ in range(20):
            is_static = calibrator.step(
                state,
                0.002,
                np.array([0.0, 0.0, 1.0]),
                np.array([0.0, 0.0, 1.0]),
                np.zeros(2),
            )

        self.assertFalse(is_static)
        self.assertEqual(state.mode, calibrator.MOVING)

    def test_acceleration_change_restarts_candidate_window(self):
        state = self._state()
        for _ in range(4):
            calibrator.step(
                state,
                0.002,
                np.array([0.0, 0.0, 1.0]),
                np.zeros(3),
                np.zeros(2),
            )

        calibrator.step(
            state,
            0.002,
            np.array([0.05, 0.0, 1.0]),
            np.zeros(3),
            np.zeros(2),
        )

        self.assertEqual(state.mode, calibrator.CANDIDATE)
        self.assertEqual(state.candidate_count, 1)
        self.assertEqual(state.candidate_elapsed_s, 0.0)

    def test_static_period_updates_running_z_bias_mean(self):
        state = self._state(enter_seconds=0.004)
        for value in [0.1, 0.1, 0.1]:
            calibrator.step(
                state,
                0.002,
                np.array([0.0, 0.0, 1.0]),
                np.array([0.0, 0.0, value]),
                np.zeros(2),
            )

        calibrator.step(
            state,
            0.002,
            np.array([0.0, 0.0, 1.0]),
            np.array([0.0, 0.0, 0.2]),
            np.zeros(2),
        )

        self.assertAlmostEqual(state.active_yaw_bias_z_dps, 0.125, places=9)

    def test_motion_exits_static_and_keeps_last_bias(self):
        state = self._state(enter_seconds=0.004)
        for _ in range(3):
            calibrator.step(
                state,
                0.002,
                np.array([0.0, 0.0, 1.0]),
                np.array([0.0, 0.0, 0.1]),
                np.zeros(2),
            )
        bias_before_motion = state.active_yaw_bias_z_dps

        is_static = calibrator.step(
            state,
            0.002,
            np.array([0.0, 0.0, 1.0]),
            np.array([0.0, 0.0, 1.0]),
            np.zeros(2),
        )

        self.assertFalse(is_static)
        self.assertEqual(state.mode, calibrator.MOVING)
        self.assertEqual(state.active_yaw_bias_z_dps, bias_before_motion)


if __name__ == "__main__":
    unittest.main()