import pytest
import numpy as np
from num_quaternions import NumQuaternion


class TestNumQuaternion:
    def test_initialization_default(self):
        """Test default quaternion initialization"""
        q = NumQuaternion()
        assert q.scalar == np.cos(np.pi / 2)
        assert q.vector.shape == (3, 1)

    def test_initialization_with_params(self):
        """Test quaternion initialization with custom parameters"""
        q = NumQuaternion(angle=np.pi/2, axis=[1.0, 0.0, 0.0])
        assert isinstance(q.scalar, (float, np.floating))
        assert q.vector.shape == (3, 1)

    def test_rotate_vector_90_degrees_x_axis(self):
        """Test rotating a vector 90 degrees around x-axis"""
        q = NumQuaternion(angle=np.pi/2, axis=[1.0, 0.0, 0.0])
        v = np.array([0.0, 1.0, 0.0])
        result = q.rotate_vector(v)

        # After 90° rotation around x-axis, [0,1,0] should become [0,0,1]
        expected = np.array([0.0, 0.0, 1.0])
        np.testing.assert_array_almost_equal(result, expected, decimal=10)

    def test_rotate_vector_180_degrees(self):
        """Test rotating a vector 180 degrees"""
        q = NumQuaternion(angle=np.pi, axis=[0.0, 0.0, 1.0])
        v = np.array([1.0, 0.0, 0.0])
        result = q.rotate_vector(v)

        # After 180° rotation around z-axis, [1,0,0] should become [-1,0,0]
        expected = np.array([-1.0, 0.0, 0.0])
        np.testing.assert_array_almost_equal(result, expected, decimal=10)

    def test_rotate_basis(self):
        """Test basis rotation"""
        q = NumQuaternion(angle=np.pi/2, axis=[1.0, 0.0, 0.0])
        v = np.array([0.0, 1.0, 0.0])
        result = q.rotate_basis(v)

        # Basis rotation should be inverse of vector rotation
        expected = np.array([0.0, 0.0, -1.0])
        np.testing.assert_array_almost_equal(result, expected, decimal=10)

    def test_get_matrix(self):
        """Test rotation matrix generation"""
        q = NumQuaternion(angle=np.pi/2, axis=[0.0, 0.0, 1.0])
        M = q.get_matrix()

        # Check matrix shape
        assert M.shape == (3, 3)

        # Check that it's a proper rotation matrix (determinant should be 1)
        det = np.linalg.det(M)
        np.testing.assert_almost_equal(det, 1.0, decimal=10)

        # Check orthogonality: M @ M.T should equal identity
        identity = np.eye(3)
        np.testing.assert_array_almost_equal(M @ M.T, identity, decimal=10)

    def test_conjugate(self):
        """Test quaternion conjugate"""
        q = NumQuaternion(angle=np.pi/4, axis=[1.0, 0.0, 0.0])
        q_conj = q._conjugate()

        # Scalar part should remain the same
        assert q_conj.scalar == q.scalar

        # Vector part should be negated
        np.testing.assert_array_almost_equal(q_conj.vector, -q.vector, decimal=10)

    def test_norma(self):
        """Test quaternion norm calculation"""
        q = NumQuaternion(angle=np.pi/2, axis=[1.0, 0.0, 0.0])
        norma = q._norma()

        # For unit quaternions (rotation quaternions), norm should be 1
        np.testing.assert_almost_equal(norma, 1.0, decimal=10)

    def test_inverse(self):
        """Test quaternion inverse"""
        q = NumQuaternion(angle=np.pi/3, axis=[1.0, 0.0, 0.0])
        q_inv = q._inv()

        # q * q_inv should give identity quaternion (scalar=1, vector=0)
        result = q._prod_quaternions(q, q_inv)
        np.testing.assert_almost_equal(result.scalar, 1.0, decimal=10)
        expected_vector = np.array([[0.0], [0.0], [0.0]])
        np.testing.assert_array_almost_equal(result.vector, expected_vector, decimal=10)

    def test_vector_product(self):
        """Test vector cross product"""
        q = NumQuaternion()
        v1 = np.array([[1.0], [0.0], [0.0]])
        v2 = np.array([[0.0], [1.0], [0.0]])
        result = q._vector_product(v1, v2)

        # Cross product of x and y should be z
        expected = np.array([[0.0], [0.0], [1.0]])
        np.testing.assert_array_almost_equal(result, expected, decimal=10)

    def test_quaternion_multiplication(self):
        """Test quaternion multiplication"""
        q1 = NumQuaternion(angle=np.pi/4, axis=[0.0, 0.0, 1.0])
        q2 = NumQuaternion(angle=np.pi/4, axis=[0.0, 0.0, 1.0])
        result = q1._prod_quaternions(q1, q2)

        # Multiplying two 45° rotations should give 90° rotation
        q_expected = NumQuaternion(angle=np.pi/2, axis=[0.0, 0.0, 1.0])
        np.testing.assert_almost_equal(result.scalar, q_expected.scalar, decimal=10)
        np.testing.assert_array_almost_equal(result.vector, q_expected.vector, decimal=10)

    def test_identity_rotation(self):
        """Test identity rotation (zero angle)"""
        q = NumQuaternion(angle=0.0, axis=[1.0, 0.0, 0.0])
        v = np.array([1.0, 2.0, 3.0])
        result = q.rotate_vector(v)

        # Identity rotation should not change the vector
        np.testing.assert_array_almost_equal(result, v, decimal=10)

    def test_full_rotation(self):
        """Test full 360 degree rotation"""
        q = NumQuaternion(angle=2*np.pi, axis=[0.0, 0.0, 1.0])
        v = np.array([1.0, 0.0, 0.0])
        result = q.rotate_vector(v)

        # 360° rotation should return the vector to its original position
        np.testing.assert_array_almost_equal(result, v, decimal=10)
