"""This module asserts correct runtime behaviour of the pyrate.act.pid module."""

# Test environment
from unittest import TestCase

# Mathematics
from numpy import array
from numpy import eye
from numpy import Inf
from numpy.linalg import norm
from numpy import vstack

# Package under test
from pyrate.act.control import Lqr


class TestLqr(TestCase):

    """Test for correct runtime behaviour in pyrate.act.control.lqr."""

    # In this context, we reproduce a common controller notation
    # pylint: disable=invalid-name

    def setUp(self) -> None:
        """Setup the LQR specification for testing."""

        # Model specification
        self.A = array([[0, 1], [0, 0]])
        self.B = array([0, 1])[:, None]
        self.C = array([1, 0])[None, :]
        self.dt = 0.5

        # Cost matrix specification
        self.Q = eye(2)
        self.R = array([[1.0]])

        # State and target
        # Target is already reached in this example
        self.state = vstack([0.0, 0.0])
        self.desired = vstack([0.0])

    def test_process_tracking(self) -> None:
        """Assert that a pandas.DataFrame with process data is created."""

        # Initialize PID controller
        lqr = Lqr(self.A, self.B, self.C, self.Q, self.R, self.dt, keep_trace=True)

        # Execute a few control steps
        state = self.state.copy()
        previous_error = Inf
        for _ in range(10):
            control_signal = lqr.control(desired=self.desired, state=self.state)
            state += self.A.dot(state) + self.B.dot(control_signal)
            error = norm(state - self.desired).item()  # Convert from numpy scalar to Python float
            assert error < previous_error or error == 0, "Error did not get lower."
            previous_error = error

        # Assert correct process tracing with LQR controller
        assert lqr.process is not None, "LQR did not keep trace of process"
        assert len(lqr.process.index) == 10, "LQR has not traced enough steps"

        # Reset PID
        lqr.reset()

        # Assert correct reset to initial state
        assert len(lqr.process.index) == 0, "LQR has not dropped process trace properly"