"""
Tests comparing SymQuaternion to NumQuaternion implementation.

This test suite verifies that SymQuaternion produces the same numerical results
as NumQuaternion when symbolic expressions are evaluated with concrete values.
"""

import numpy as np
import sympy as sp
import pytest

from sym_quaternions.quaternion import SymQuaternion
from num_quaternions.quaternion import NumQuaternion


class TestQuaternionComparison:
    """Compare SymQuaternion to NumQuaternion for various operations."""

    def assert_quaternions_equal(self, sym_q, num_q, tolerance=1e-10):
        """
        Assert that symbolic and numeric quaternions are equal.

        Args:
            sym_q: SymQuaternion instance
            num_q: NumQuaternion instance
            tolerance: numerical tolerance for comparison
        """
        assert abs(float(sym_q.w) - float(num_q.scalar)) < tolerance, \
            f"w/scalar mismatch: sym={sym_q.w}, num={num_q.scalar}"
        assert abs(float(sym_q.x) - float(num_q.vector[0, 0])) < tolerance, \
            f"x mismatch: sym={sym_q.x}, num={num_q.vector[0, 0]}"
        assert abs(float(sym_q.y) - float(num_q.vector[1, 0])) < tolerance, \
            f"y mismatch: sym={sym_q.y}, num={num_q.vector[1, 0]}"
        assert abs(float(sym_q.z) - float(num_q.vector[2, 0])) < tolerance, \
            f"z mismatch: sym={sym_q.z}, num={num_q.vector[2, 0]}"

    def assert_matrices_equal(self, sym_mat, num_mat, tolerance=1e-10):
        """
        Assert that symbolic and numeric matrices are equal.

        Args:
            sym_mat: sympy.Matrix
            num_mat: numpy array
            tolerance: numerical tolerance for comparison
        """
        for i in range(3):
            for j in range(3):
                sym_val = float(sym_mat[i, j])
                num_val = float(num_mat[i, j])
                assert abs(sym_val - num_val) < tolerance, \
                    f"Matrix element [{i},{j}] mismatch: sym={sym_val}, num={num_val}"

    def assert_vectors_equal(self, sym_vec, num_vec, tolerance=1e-10):
        """
        Assert that symbolic and numeric vectors are equal.

        Args:
            sym_vec: sympy.Matrix or numpy array
            num_vec: numpy array
            tolerance: numerical tolerance for comparison
        """
        for i in range(3):
            sym_val = float(sym_vec[i])
            num_val = float(num_vec[i])
            assert abs(sym_val - num_val) < tolerance, \
                f"Vector element [{i}] mismatch: sym={sym_val}, num={num_val}"


class TestIdentityQuaternion(TestQuaternionComparison):
    """Test identity quaternion (no rotation)."""

    def test_identity_creation(self):
        """Test identity quaternion creation."""
        sym_q = SymQuaternion()
        num_q = NumQuaternion(angle=0.0, axis=[0.0, 0.0, 1.0])

        self.assert_quaternions_equal(sym_q, num_q)

    def test_identity_rotation_matrix(self):
        """Test identity quaternion produces identity rotation matrix."""
        sym_q = SymQuaternion()
        num_q = NumQuaternion(angle=0.0, axis=[0.0, 0.0, 1.0])

        sym_mat = sym_q.to_rotation_matrix()
        num_mat = num_q.get_matrix()

        self.assert_matrices_equal(sym_mat, num_mat)

        # Check it's actually identity
        identity = np.eye(3)
        self.assert_matrices_equal(sym_mat, identity)

    def test_identity_vector_rotation(self):
        """Test identity quaternion doesn't change vectors."""
        sym_q = SymQuaternion()
        num_q = NumQuaternion(angle=0.0, axis=[0.0, 0.0, 1.0])

        test_vectors = [
            [1.0, 0.0, 0.0],
            [0.0, 1.0, 0.0],
            [0.0, 0.0, 1.0],
            [1.0, 2.0, 3.0],
        ]

        for vec in test_vectors:
            sym_result = sym_q.rotate_vector(sp.Matrix(vec))
            num_result = num_q.rotate_vector(np.array(vec))

            self.assert_vectors_equal(sym_result, num_result)
            self.assert_vectors_equal(sym_result, vec)


class TestAxisAlignedRotations(TestQuaternionComparison):
    """Test rotations around principal axes."""

    @pytest.mark.parametrize("angle,axis", [
        (np.pi/2, [1.0, 0.0, 0.0]),  # 90° around X
        (np.pi, [1.0, 0.0, 0.0]),    # 180° around X
        (np.pi/2, [0.0, 1.0, 0.0]),  # 90° around Y
        (np.pi, [0.0, 1.0, 0.0]),    # 180° around Y
        (np.pi/2, [0.0, 0.0, 1.0]),  # 90° around Z
        (np.pi, [0.0, 0.0, 1.0]),    # 180° around Z
    ])
    def test_axis_rotation_creation(self, angle, axis):
        """Test quaternion creation from angle-axis."""
        sym_q = SymQuaternion(angle=angle, axis=axis)
        num_q = NumQuaternion(angle=angle, axis=axis)

        self.assert_quaternions_equal(sym_q, num_q)

    @pytest.mark.parametrize("angle,axis", [
        (np.pi/2, [1.0, 0.0, 0.0]),
        (np.pi, [1.0, 0.0, 0.0]),
        (np.pi/2, [0.0, 1.0, 0.0]),
        (np.pi, [0.0, 1.0, 0.0]),
        (np.pi/2, [0.0, 0.0, 1.0]),
        (np.pi, [0.0, 0.0, 1.0]),
    ])
    def test_axis_rotation_matrix(self, angle, axis):
        """Test rotation matrices for axis-aligned rotations."""
        sym_q = SymQuaternion(angle=angle, axis=axis)
        num_q = NumQuaternion(angle=angle, axis=axis)

        sym_mat = sym_q.to_rotation_matrix()
        num_mat = num_q.get_matrix()

        self.assert_matrices_equal(sym_mat, num_mat)

    def test_90deg_x_rotation_vector(self):
        """Test 90° rotation around X axis."""
        angle = np.pi / 2
        axis = [1.0, 0.0, 0.0]

        sym_q = SymQuaternion(angle=angle, axis=axis)
        num_q = NumQuaternion(angle=angle, axis=axis)

        # Rotate [0, 1, 0] -> should give [0, 0, 1]
        vec = [0.0, 1.0, 0.0]
        sym_result = sym_q.rotate_vector(sp.Matrix(vec))
        num_result = num_q.rotate_vector(np.array(vec))

        self.assert_vectors_equal(sym_result, num_result)
        # Check expected result (with tolerance for numerical precision)
        expected = [0.0, 0.0, 1.0]
        self.assert_vectors_equal(sym_result, expected)

    def test_90deg_z_rotation_vector(self):
        """Test 90° rotation around Z axis."""
        angle = np.pi / 2
        axis = [0.0, 0.0, 1.0]

        sym_q = SymQuaternion(angle=angle, axis=axis)
        num_q = NumQuaternion(angle=angle, axis=axis)

        # Rotate [1, 0, 0] -> should give [0, 1, 0]
        vec = [1.0, 0.0, 0.0]
        sym_result = sym_q.rotate_vector(sp.Matrix(vec))
        num_result = num_q.rotate_vector(np.array(vec))

        self.assert_vectors_equal(sym_result, num_result)
        expected = [0.0, 1.0, 0.0]
        self.assert_vectors_equal(sym_result, expected)


class TestSymbolicEvaluation(TestQuaternionComparison):
    """Test symbolic quaternions with parameter substitution."""

    def test_symbolic_angle_substitution(self):
        """Test symbolic angle that gets substituted."""
        theta = sp.Symbol('theta', real=True)
        axis = [0.0, 0.0, 1.0]

        # Create symbolic quaternion
        sym_q = SymQuaternion(angle=theta, axis=axis)

        # Substitute concrete value
        theta_val = np.pi / 3  # 60 degrees
        sym_q_eval = sym_q.eval({theta: theta_val})

        # Create numeric quaternion with same value
        num_q = NumQuaternion(angle=theta_val, axis=axis)

        self.assert_quaternions_equal(sym_q_eval, num_q)

    def test_symbolic_axis_substitution(self):
        """Test symbolic axis components."""
        angle = np.pi / 4
        ax, ay, az = sp.symbols('ax ay az', real=True)

        # Create symbolic quaternion
        sym_q = SymQuaternion(angle=angle, axis=[ax, ay, az])

        # Substitute concrete values (normalized)
        axis_vals = np.array([1.0, 1.0, 1.0])
        axis_vals = axis_vals / np.linalg.norm(axis_vals)

        sym_q_eval = sym_q.eval({ax: axis_vals[0], ay: axis_vals[1], az: axis_vals[2]})

        # Create numeric quaternion
        num_q = NumQuaternion(angle=angle, axis=axis_vals.tolist())

        self.assert_quaternions_equal(sym_q_eval, num_q)

    def test_symbolic_rotation_matrix(self):
        """Test rotation matrix with symbolic parameters."""
        theta = sp.Symbol('theta', real=True)
        sym_q = SymQuaternion(angle=theta, axis=[0.0, 0.0, 1.0])

        # Get symbolic rotation matrix
        sym_mat = sym_q.to_rotation_matrix()

        # Substitute value
        theta_val = np.pi / 6  # 30 degrees
        sym_mat_eval = sym_mat.subs(theta, theta_val)

        # Get numeric rotation matrix
        num_q = NumQuaternion(angle=theta_val, axis=[0.0, 0.0, 1.0])
        num_mat = num_q.get_matrix()

        self.assert_matrices_equal(sym_mat_eval, num_mat)

    def test_symbolic_vector_rotation(self):
        """Test rotating a vector with symbolic quaternion."""
        theta = sp.Symbol('theta', real=True)
        sym_q = SymQuaternion(angle=theta, axis=[1.0, 0.0, 0.0])

        # Rotate symbolic vector
        vec = sp.Matrix([0, 1, 0])
        sym_result = sym_q.rotate_vector(vec)

        # Substitute value
        theta_val = np.pi / 2
        sym_result_eval = sym_result.subs(theta, theta_val)

        # Compare with numeric
        num_q = NumQuaternion(angle=theta_val, axis=[1.0, 0.0, 0.0])
        num_result = num_q.rotate_vector(np.array([0.0, 1.0, 0.0]))

        self.assert_vectors_equal(sym_result_eval, num_result)


class TestQuaternionOperations(TestQuaternionComparison):
    """Test quaternion operations (conjugate, inverse, normalize, multiply)."""

    def test_conjugate(self):
        """Test quaternion conjugate."""
        angle = np.pi / 3
        axis = [1.0, 2.0, 3.0]

        sym_q = SymQuaternion(angle=angle, axis=axis)
        num_q = NumQuaternion(angle=angle, axis=axis)

        sym_conj = sym_q.conjugate()
        num_conj = num_q._conjugate()

        self.assert_quaternions_equal(sym_conj, num_conj)

    def test_norm(self):
        """Test quaternion norm."""
        angle = np.pi / 4
        axis = [1.0, 0.0, 0.0]

        sym_q = SymQuaternion(angle=angle, axis=axis)
        num_q = NumQuaternion(angle=angle, axis=axis)

        sym_norm = float(sym_q.norm())
        num_norm = float(num_q._norma())

        assert abs(sym_norm - num_norm) < 1e-10

    def test_inverse(self):
        """Test quaternion inverse."""
        angle = np.pi / 5
        axis = [1.0, 1.0, 0.0]

        sym_q = SymQuaternion(angle=angle, axis=axis)
        num_q = NumQuaternion(angle=angle, axis=axis)

        sym_inv = sym_q.inv()
        num_inv = num_q._inv()

        self.assert_quaternions_equal(sym_inv, num_inv)

    def test_normalize(self):
        """Test quaternion normalization."""
        # Create non-normalized quaternion
        sym_q = SymQuaternion(w=2.0, x=1.0, y=1.0, z=1.0)

        # Normalize
        sym_norm = sym_q.normalize()

        # Check norm is 1
        norm_squared = float(sym_norm.norm())
        assert abs(norm_squared - 1.0) < 1e-10

    def test_rotation_basis(self):
        """Test basis rotation (inverse of vector rotation)."""
        angle = np.pi / 4
        axis = [0.0, 1.0, 0.0]

        sym_q = SymQuaternion(angle=angle, axis=axis)
        num_q = NumQuaternion(angle=angle, axis=axis)

        vec = [1.0, 0.0, 1.0]
        sym_result = sym_q.rotate_basis(sp.Matrix(vec))
        num_result = num_q.rotate_basis(np.array(vec))

        self.assert_vectors_equal(sym_result, num_result)


class TestRandomParameters(TestQuaternionComparison):
    """Test with random varied parameters."""

    @pytest.mark.parametrize("seed", range(10))
    def test_random_angle_axis(self, seed):
        """Test random angle-axis rotations."""
        np.random.seed(seed)

        # Random angle between 0 and 2π
        angle = np.random.uniform(0, 2 * np.pi)

        # Random axis (will be normalized)
        axis = np.random.randn(3)
        axis = axis / np.linalg.norm(axis)

        sym_q = SymQuaternion(angle=angle, axis=axis.tolist())
        num_q = NumQuaternion(angle=angle, axis=axis.tolist())

        self.assert_quaternions_equal(sym_q, num_q)

    @pytest.mark.parametrize("seed", range(10))
    def test_random_rotation_matrix(self, seed):
        """Test rotation matrices with random parameters."""
        np.random.seed(seed)

        angle = np.random.uniform(0, 2 * np.pi)
        axis = np.random.randn(3)
        axis = axis / np.linalg.norm(axis)

        sym_q = SymQuaternion(angle=angle, axis=axis.tolist())
        num_q = NumQuaternion(angle=angle, axis=axis.tolist())

        sym_mat = sym_q.to_rotation_matrix()
        num_mat = num_q.get_matrix()

        self.assert_matrices_equal(sym_mat, num_mat)

    @pytest.mark.parametrize("seed", range(10))
    def test_random_vector_rotation(self, seed):
        """Test vector rotations with random parameters."""
        np.random.seed(seed)

        angle = np.random.uniform(0, 2 * np.pi)
        axis = np.random.randn(3)
        axis = axis / np.linalg.norm(axis)

        vec = np.random.randn(3)

        sym_q = SymQuaternion(angle=angle, axis=axis.tolist())
        num_q = NumQuaternion(angle=angle, axis=axis.tolist())

        sym_result = sym_q.rotate_vector(sp.Matrix(vec.tolist()))
        num_result = num_q.rotate_vector(vec)

        self.assert_vectors_equal(sym_result, num_result)

    @pytest.mark.parametrize("seed", range(5))
    def test_random_basis_rotation(self, seed):
        """Test basis rotations with random parameters."""
        np.random.seed(seed)

        angle = np.random.uniform(0, 2 * np.pi)
        axis = np.random.randn(3)
        axis = axis / np.linalg.norm(axis)

        vec = np.random.randn(3)

        sym_q = SymQuaternion(angle=angle, axis=axis.tolist())
        num_q = NumQuaternion(angle=angle, axis=axis.tolist())

        sym_result = sym_q.rotate_basis(sp.Matrix(vec.tolist()))
        num_result = num_q.rotate_basis(vec)

        self.assert_vectors_equal(sym_result, num_result)

    @pytest.mark.parametrize("seed", range(5))
    def test_random_operations_chain(self, seed):
        """Test chained operations with random parameters."""
        np.random.seed(seed)

        angle = np.random.uniform(0, np.pi)
        axis = np.random.randn(3)
        axis = axis / np.linalg.norm(axis)

        sym_q = SymQuaternion(angle=angle, axis=axis.tolist())
        num_q = NumQuaternion(angle=angle, axis=axis.tolist())

        # Test conjugate -> inverse -> normalize chain
        sym_result = sym_q.conjugate().inv().normalize()
        num_result = num_q._conjugate()._inv()

        # Normalize num_result manually
        num_norm = num_result._norma()
        num_result.scalar = num_result.scalar / np.sqrt(num_norm)
        num_result.vector = num_result.vector / np.sqrt(num_norm)

        self.assert_quaternions_equal(sym_result, num_result, tolerance=1e-9)


class TestEdgeCases(TestQuaternionComparison):
    """Test edge cases and special values."""

    def test_zero_angle_rotation(self):
        """Test rotation with zero angle (identity)."""
        sym_q = SymQuaternion(angle=0.0, axis=[1.0, 0.0, 0.0])
        num_q = NumQuaternion(angle=0.0, axis=[1.0, 0.0, 0.0])

        self.assert_quaternions_equal(sym_q, num_q)

        # Should be identity quaternion
        assert abs(float(sym_q.w) - 1.0) < 1e-10
        assert abs(float(sym_q.x)) < 1e-10
        assert abs(float(sym_q.y)) < 1e-10
        assert abs(float(sym_q.z)) < 1e-10

    def test_full_rotation(self):
        """Test 360° rotation (should be identity)."""
        angle = 2 * np.pi
        axis = [0.0, 0.0, 1.0]

        sym_q = SymQuaternion(angle=angle, axis=axis)
        num_q = NumQuaternion(angle=angle, axis=axis)

        self.assert_quaternions_equal(sym_q, num_q)

        # Rotate a vector - should return to original
        vec = [1.0, 2.0, 3.0]
        sym_result = sym_q.rotate_vector(sp.Matrix(vec))
        num_result = num_q.rotate_vector(np.array(vec))

        self.assert_vectors_equal(sym_result, num_result)
        self.assert_vectors_equal(sym_result, vec)

    def test_opposite_rotations(self):
        """Test that opposite rotations cancel out."""
        angle = np.pi / 3
        axis = [1.0, 1.0, 1.0]

        sym_q1 = SymQuaternion(angle=angle, axis=axis)
        sym_q2 = SymQuaternion(angle=-angle, axis=axis)

        # Rotating by q1 then q2 should give identity
        vec = sp.Matrix([1.0, 2.0, 3.0])
        result = sym_q2.rotate_vector(sym_q1.rotate_vector(vec))

        self.assert_vectors_equal(result, [1.0, 2.0, 3.0])


if __name__ == "__main__":
    pytest.main([__file__, "-v"])
