import numpy as np
import pytest

from navlib.math.quaternion import (
    left_qmatrix,
    quat_conj,
    quat_inv,
    quat_normalize,
    quat_positive_scalar,
    right_qmatrix,
)


def test_left_qmatrix():
    """Test the left_qmatrix function."""
    # Test case 1: Identity quaternion
    q_identity = np.array([0, 0, 0, 1])
    expected_left = np.array([[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]])
    np.testing.assert_array_equal(left_qmatrix(q_identity), expected_left)
    np.testing.assert_array_equal(left_qmatrix(q_identity.tolist()), expected_left)

    # Test case 7: 2D array that becomes 1D after squeeze (should work)
    q_2d = np.array([[0, 0, 0, 1]])
    np.testing.assert_array_equal(left_qmatrix(q_2d), expected_left)

    # Test case 2: 90-degree rotation around Z-axis
    q_z90 = np.array([0, 0, np.sqrt(2) / 2, np.sqrt(2) / 2])
    expected_left_z90 = np.array(
        [
            [np.sqrt(2) / 2, -np.sqrt(2) / 2, 0, 0],
            [np.sqrt(2) / 2, np.sqrt(2) / 2, 0, 0],
            [0, 0, np.sqrt(2) / 2, np.sqrt(2) / 2],
            [0, 0, -np.sqrt(2) / 2, np.sqrt(2) / 2],
        ]
    )
    np.testing.assert_array_almost_equal(left_qmatrix(q_z90), expected_left_z90)

    # Test case 3: General normalized quaternion
    q = np.array([0.1, 0.2, 0.3, np.sqrt(1 - 0.1**2 - 0.2**2 - 0.3**2)])
    w, x, y, z = q[3], q[0], q[1], q[2]
    expected_left = np.array(
        [[w, -z, y, x], [z, w, -x, y], [-y, x, w, z], [-x, -y, -z, w]]
    )
    np.testing.assert_array_almost_equal(left_qmatrix(q), expected_left)

    # Test case 4: Wrong input type
    with pytest.raises(TypeError):
        left_qmatrix("not a numpy array")

    # Test case 5: Wrong number of elements
    with pytest.raises(ValueError):
        left_qmatrix(np.array([1, 2, 3]))

    # Test case 6: Non-normalized quaternion
    with pytest.raises(ValueError):
        left_qmatrix(np.array([1, 2, 3, 4]))


def test_right_qmatrix():
    """Test the right_qmatrix function."""
    # Test case 1: Identity quaternion
    q_identity = np.array([0, 0, 0, 1])
    expected_right = np.array([[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]])
    np.testing.assert_array_equal(right_qmatrix(q_identity), expected_right)
    np.testing.assert_array_equal(right_qmatrix(q_identity.tolist()), expected_right)

    # Test case 2: 90-degree rotation around Z-axis
    q_z90 = np.array([0, 0, np.sqrt(2) / 2, np.sqrt(2) / 2])
    expected_right_z90 = np.array(
        [
            [np.sqrt(2) / 2, np.sqrt(2) / 2, 0, 0],
            [-np.sqrt(2) / 2, np.sqrt(2) / 2, 0, 0],
            [0, 0, np.sqrt(2) / 2, np.sqrt(2) / 2],
            [0, 0, -np.sqrt(2) / 2, np.sqrt(2) / 2],
        ]
    )
    np.testing.assert_array_almost_equal(right_qmatrix(q_z90), expected_right_z90)

    # Test case 3: General normalized quaternion
    q = np.array([0.1, 0.2, 0.3, np.sqrt(1 - 0.1**2 - 0.2**2 - 0.3**2)])
    w, x, y, z = q[3], q[0], q[1], q[2]
    expected_right = np.array(
        [[w, z, -y, x], [-z, w, x, y], [y, -x, w, z], [-x, -y, -z, w]]
    )
    np.testing.assert_array_almost_equal(right_qmatrix(q), expected_right)

    # Test case 4: Verify quaternion multiplication property
    q1 = np.array([0.1, 0.2, 0.3, np.sqrt(1 - 0.1**2 - 0.2**2 - 0.3**2)])
    q2 = np.array([0.2, 0.1, 0.4, np.sqrt(1 - 0.2**2 - 0.1**2 - 0.4**2)])

    # q1 * q2 = L(q1) * q2 = R(q2) * q1
    left_result = left_qmatrix(q1) @ q2
    right_result = right_qmatrix(q2) @ q1
    np.testing.assert_array_almost_equal(left_result, right_result)

    # Test case 5: Wrong input type
    with pytest.raises(TypeError):
        right_qmatrix("not a numpy array")

    # Test case 6: Wrong number of elements
    with pytest.raises(ValueError):
        right_qmatrix(np.array([1, 2, 3]))

    # Test case 7: Non-normalized quaternion
    with pytest.raises(ValueError):
        right_qmatrix(np.array([1, 2, 3, 4]))


def test_quat_conj():
    """Test the quat_conj function."""
    # Test case 1: Identity quaternion
    q_identity = np.array([0, 0, 0, 1])
    expected_conj = np.array([0, 0, 0, 1])
    np.testing.assert_array_equal(quat_conj(q_identity), expected_conj)
    np.testing.assert_array_equal(quat_conj(q_identity.tolist()), expected_conj)

    # Test case 2: General quaternion
    q = np.array([0.1, 0.2, 0.3, 0.4])
    expected_conj = np.array([-0.1, -0.2, -0.3, 0.4])
    np.testing.assert_array_equal(quat_conj(q), expected_conj)

    # Test case 3: Pure quaternion (w = 0)
    q_pure = np.array([1, 2, 3, 0])
    expected_conj_pure = np.array([-1, -2, -3, 0])
    np.testing.assert_array_equal(quat_conj(q_pure), expected_conj_pure)

    # Test case 4: Conjugate of conjugate should equal original
    q = np.array([0.1, 0.2, 0.3, 0.4])
    np.testing.assert_array_equal(quat_conj(quat_conj(q)), q)

    # Test case 5: Wrong input type
    with pytest.raises(TypeError):
        quat_conj("not a numpy array")

    # Test case 6: Wrong number of elements
    with pytest.raises(ValueError):
        quat_conj(np.array([1, 2, 3]))

    # Test case 7: 2D array that becomes 1D after squeeze (should work)
    q_2d = np.array([[0.1, 0.2, 0.3, 0.4]])
    expected = np.array([-0.1, -0.2, -0.3, 0.4])
    np.testing.assert_array_equal(quat_conj(q_2d), expected)


def test_quat_inv():
    """Test the quat_inv function."""
    # Test case 1: Identity quaternion
    q_identity = np.array([0, 0, 0, 1])
    expected_inv = np.array([0, 0, 0, 1])
    np.testing.assert_array_equal(quat_inv(q_identity), expected_inv)
    np.testing.assert_array_equal(quat_inv(q_identity.tolist()), expected_inv)

    # Test case 2: Unit quaternion
    q_unit = np.array([0.1, 0.2, 0.3, np.sqrt(1 - 0.1**2 - 0.2**2 - 0.3**2)])
    q_inv = quat_inv(q_unit)
    # For unit quaternions, inverse equals conjugate
    np.testing.assert_array_almost_equal(q_inv, quat_conj(q_unit))

    # Test case 3: General quaternion
    q = np.array([0.1, 0.2, 0.3, 0.4])
    q_inv = quat_inv(q)
    expected_inv = quat_conj(q) / np.dot(q, q)
    np.testing.assert_array_almost_equal(q_inv, expected_inv)

    # Test case 4: Verify quaternion property q * q^(-1) has magnitude ||q||^2
    q = np.array([0.1, 0.2, 0.3, 0.4])
    q_inv = quat_inv(q)
    # Normalize q for left_qmatrix
    q_norm = quat_normalize(q)
    q_inv_norm = quat_normalize(q_inv)
    # Manual quaternion multiplication using normalized quaternions
    result = left_qmatrix(q_norm) @ q_inv_norm
    # Result should be close to [0, 0, 0, 1] for normalized quaternions
    expected_norm = np.array([0, 0, 0, 1])
    np.testing.assert_array_almost_equal(result, expected_norm, decimal=5)

    # Test case 5: Wrong input type
    with pytest.raises(TypeError):
        quat_inv("not a numpy array")

    # Test case 6: Wrong number of elements
    with pytest.raises(ValueError):
        quat_inv(np.array([1, 2, 3]))


def test_quat_normalize():
    """Test the quat_normalize function."""
    # Test case 1: Already normalized quaternion
    q_unit = np.array([0, 0, 0, 1])
    np.testing.assert_array_almost_equal(quat_normalize(q_unit), q_unit)
    np.testing.assert_array_almost_equal(quat_normalize(q_unit.tolist()), q_unit)

    # Test case 2: Non-normalized quaternion
    q = np.array([1, 2, 3, 4])
    q_norm = quat_normalize(q)
    expected_norm = q / np.linalg.norm(q)
    np.testing.assert_array_almost_equal(q_norm, expected_norm)

    # Verify the result is normalized
    assert np.isclose(np.linalg.norm(q_norm), 1.0)

    # Test case 3: Small quaternion
    q_small = np.array([0.1, 0.1, 0.1, 0.1])
    q_norm = quat_normalize(q_small)
    assert np.isclose(np.linalg.norm(q_norm), 1.0)

    # Test case 4: Zero quaternion (should raise error)
    with pytest.raises(ValueError):
        quat_normalize(np.array([0, 0, 0, 0]))

    # Test case 5: Very small quaternion (should raise error)
    with pytest.raises(ValueError):
        quat_normalize(np.array([1e-10, 1e-10, 1e-10, 1e-10]))

    # Test case 6: Wrong input type
    with pytest.raises(TypeError):
        quat_normalize("not a numpy array")

    # Test case 7: Wrong number of elements
    with pytest.raises(ValueError):
        quat_normalize(np.array([1, 2, 3]))

    # Test case 8: 2D array that becomes 1D after squeeze (should work)
    q_2d = np.array([[1, 2, 3, 4]])
    q_norm = quat_normalize(q_2d)
    assert np.isclose(np.linalg.norm(q_norm), 1.0)


def test_quat_positive_scalar():
    """Test the quat_positive_scalar function."""
    # Test case 1: Already positive scalar
    q_pos = np.array([0.1, 0.2, 0.3, 0.5])
    np.testing.assert_array_equal(quat_positive_scalar(q_pos), q_pos)
    np.testing.assert_array_equal(quat_positive_scalar(q_pos.tolist()), q_pos)

    # Test case 2: Negative scalar
    q_neg = np.array([0.1, 0.2, 0.3, -0.5])
    expected = np.array([-0.1, -0.2, -0.3, 0.5])
    np.testing.assert_array_equal(quat_positive_scalar(q_neg), expected)

    # Test case 3: Zero scalar
    q_zero = np.array([0.1, 0.2, 0.3, 0.0])
    np.testing.assert_array_equal(quat_positive_scalar(q_zero), q_zero)

    # Test case 4: Identity quaternion
    q_identity = np.array([0, 0, 0, 1])
    np.testing.assert_array_equal(quat_positive_scalar(q_identity), q_identity)

    # Test case 5: Negative identity quaternion
    q_neg_identity = np.array([0, 0, 0, -1])
    expected_identity = np.array([0, 0, 0, 1])
    np.testing.assert_array_equal(
        quat_positive_scalar(q_neg_identity), expected_identity
    )

    # Test case 6: Verify double application gives same result
    q = np.array([0.1, 0.2, 0.3, -0.5])
    result1 = quat_positive_scalar(q)
    result2 = quat_positive_scalar(result1)
    np.testing.assert_array_equal(result1, result2)

    # Test case 7: Wrong input type
    with pytest.raises(TypeError):
        quat_positive_scalar("not a numpy array")

    # Test case 8: Wrong number of elements
    with pytest.raises(ValueError):
        quat_positive_scalar(np.array([1, 2, 3]))

    # Test case 9: 2D array that becomes 1D after squeeze (should work)
    q_2d = np.array([[0.1, 0.2, 0.3, -0.5]])
    expected = np.array([-0.1, -0.2, -0.3, 0.5])
    np.testing.assert_array_equal(quat_positive_scalar(q_2d), expected)


def test_quaternion_mathematical_properties():
    """Test mathematical properties of quaternion operations."""
    # Test quaternion multiplication identity: q * I = I * q = q
    q = quat_normalize(np.array([0.1, 0.2, 0.3, 0.4]))
    q_identity = np.array([0, 0, 0, 1])

    # Test left multiplication with identity
    result_left = left_qmatrix(q_identity) @ q
    np.testing.assert_array_almost_equal(result_left, q)

    # Test right multiplication with identity
    result_right = right_qmatrix(q_identity) @ q
    np.testing.assert_array_almost_equal(result_right, q)

    # Test quaternion inverse property: q * q^(-1) = identity (for unit quaternions)
    q_unit = quat_normalize(np.array([0.1, 0.2, 0.3, 0.4]))
    q_inv = quat_inv(q_unit)
    q_inv_norm = quat_normalize(q_inv)  # Normalize for matrix multiplication

    result = left_qmatrix(q_unit) @ q_inv_norm
    expected = np.array([0, 0, 0, 1])
    np.testing.assert_array_almost_equal(result, expected, decimal=10)

    # Test conjugate property: ||q*|| = ||q||
    q = np.array([0.1, 0.2, 0.3, 0.4])
    q_conj = quat_conj(q)
    np.testing.assert_almost_equal(np.linalg.norm(q), np.linalg.norm(q_conj))

    # Test positive scalar consistency
    q = quat_normalize(np.array([0.1, 0.2, 0.3, -0.4]))
    q_pos = quat_positive_scalar(q)
    assert q_pos[3] >= 0.0
    np.testing.assert_almost_equal(np.linalg.norm(q), np.linalg.norm(q_pos))
