# Path: robel/utils/math_utils.py
# def average_quaternions(quaternions: Sequence[np.ndarray]) -> np.ndarray:
#     """Returns the average of the given quaternions.
# 
#     Args:
#         quaternions: A list of quaternions to average.
# 
#     Returns:
#         The averaged quaternion.
#     """
#     # Implements the algorithm from:
#     # Markley, F. L., Cheng, Y., Crassidis, J. L., & Oshman, Y. (2007).
#     # Averaging quaternions. Journal of Guidance, Control, and Dynamics,
#     # 30(4), 1193-1197.
#     n_quat = len(quaternions)
#     assert n_quat > 0, 'Must provide at least one quaternion.'
#     weight = 1.0 / n_quat  # Uniform weighting for all quaternions.
#     q_matrix = np.vstack(quaternions)
#     assert q_matrix.shape == (n_quat, 4)
#     m_matrix = np.matmul(weight * np.transpose(q_matrix), q_matrix)
#     _, eig_vecs = np.linalg.eigh(m_matrix)
#     # The final eigenvector corresponds to the largest eigenvalue.
#     return eig_vecs[:, -1]
# 
# def calculate_cosine(vec1: np.ndarray, vec2: np.ndarray) -> np.ndarray:
#     """Calculates the cosine angle between two vectors.
# 
#     This computes cos(theta) = dot(v1, v2) / (norm(v1) * norm(v2))
# 
#     Args:
#         vec1: The first vector. This can have a batch dimension.
#         vec2: The second vector. This can have a batch dimension.
# 
#     Returns:
#         The cosine angle between the two vectors, with the same batch dimension
#         as the given vectors.
#     """
#     if np.shape(vec1) != np.shape(vec2):
#         raise ValueError('{} must have the same shape as {}'.format(vec1, vec2))
#     ndim = np.ndim(vec1)
#     if ndim < 1 or ndim > 2:
#         raise ValueError('{} must be 1 or 2 dimensions'.format(vec1))
#     axis = 1 if ndim == 2 else 0
#     norm_product = (
#         np.linalg.norm(vec1, axis=axis) * np.linalg.norm(vec2, axis=axis))
#     zero_norms = norm_product == 0
#     if np.any(zero_norms):
#         logging.warning(
#             '%s or %s is all 0s; this may be normal during initialization.',
#             str(vec1), str(vec2))
#         if ndim == 2:
#             norm_product[zero_norms] = 1
#         else:
#             norm_product = 1
#     # Return the batched dot product.
#     return np.einsum('...i,...i', vec1, vec2) / norm_product


from absl.testing import absltest
from transforms3d.euler import euler2quat, quat2euler
from robel.utils.math_utils import average_quaternions, calculate_cosine
import numpy as np
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Tests for math_utils."""




class AverageQuaternionsTest(absltest.TestCase):
    """Tests for `average_quaternions`."""

    def test_identity(self):
        """Average one quaternion should equal itself."""
        test_quat = euler2quat(np.pi / 4, np.pi / 4, np.pi / 4)
        avg_quat = average_quaternions([test_quat])
        np.testing.assert_array_almost_equal(avg_quat, test_quat)

    def test_multiple_identity(self):
        """Average multiple copies of a quaternion should equal itself."""
        test_quat = euler2quat(np.pi / 4, np.pi / 4, np.pi / 4)
        avg_quat = average_quaternions([test_quat, test_quat, test_quat])
        np.testing.assert_array_almost_equal(avg_quat, test_quat)

    def test_average_two(self):
        """Averaging two different quaternions."""
        quat1 = euler2quat(np.pi / 4, 0, 0)
        quat2 = euler2quat(-np.pi / 4, 0, 0)

# Path: robel/utils/math_utils.py
# def average_quaternions(quaternions: Sequence[np.ndarray]) -> np.ndarray:
#     """Returns the average of the given quaternions.
# 
#     Args:
#         quaternions: A list of quaternions to average.
# 
#     Returns:
#         The averaged quaternion.
#     """
#     # Implements the algorithm from:
#     # Markley, F. L., Cheng, Y., Crassidis, J. L., & Oshman, Y. (2007).
#     # Averaging quaternions. Journal of Guidance, Control, and Dynamics,
#     # 30(4), 1193-1197.
#     n_quat = len(quaternions)
#     assert n_quat > 0, 'Must provide at least one quaternion.'
#     weight = 1.0 / n_quat  # Uniform weighting for all quaternions.
#     q_matrix = np.vstack(quaternions)
#     assert q_matrix.shape == (n_quat, 4)
#     m_matrix = np.matmul(weight * np.transpose(q_matrix), q_matrix)
#     _, eig_vecs = np.linalg.eigh(m_matrix)
#     # The final eigenvector corresponds to the largest eigenvalue.
#     return eig_vecs[:, -1]
# 
# def calculate_cosine(vec1: np.ndarray, vec2: np.ndarray) -> np.ndarray:
#     """Calculates the cosine angle between two vectors.
# 
#     This computes cos(theta) = dot(v1, v2) / (norm(v1) * norm(v2))
# 
#     Args:
#         vec1: The first vector. This can have a batch dimension.
#         vec2: The second vector. This can have a batch dimension.
# 
#     Returns:
#         The cosine angle between the two vectors, with the same batch dimension
#         as the given vectors.
#     """
#     if np.shape(vec1) != np.shape(vec2):
#         raise ValueError('{} must have the same shape as {}'.format(vec1, vec2))
#     ndim = np.ndim(vec1)
#     if ndim < 1 or ndim > 2:
#         raise ValueError('{} must be 1 or 2 dimensions'.format(vec1))
#     axis = 1 if ndim == 2 else 0
#     norm_product = (
#         np.linalg.norm(vec1, axis=axis) * np.linalg.norm(vec2, axis=axis))
#     zero_norms = norm_product == 0
#     if np.any(zero_norms):
#         logging.warning(
#             '%s or %s is all 0s; this may be normal during initialization.',
#             str(vec1), str(vec2))
#         if ndim == 2:
#             norm_product[zero_norms] = 1
#         else:
#             norm_product = 1
#     # Return the batched dot product.
#     return np.einsum('...i,...i', vec1, vec2) / norm_product


# Path: robel/utils/math_utils_test.py
from absl.testing import absltest
from transforms3d.euler import euler2quat, quat2euler
from robel.utils.math_utils import average_quaternions, calculate_cosine
import numpy as np

# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Tests for math_utils."""




class AverageQuaternionsTest(absltest.TestCase):
    """Tests for `average_quaternions`."""

    def test_identity(self):
        """Average one quaternion should equal itself."""
        test_quat = euler2quat(np.pi / 4, np.pi / 4, np.pi / 4)
        avg_quat = average_quaternions([test_quat])
        np.testing.assert_array_almost_equal(avg_quat, test_quat)

    def test_multiple_identity(self):
        """Average multiple copies of a quaternion should equal itself."""
        test_quat = euler2quat(np.pi / 4, np.pi / 4, np.pi / 4)
        avg_quat = average_quaternions([test_quat, test_quat, test_quat])
        np.testing.assert_array_almost_equal(avg_quat, test_quat)

    def test_average_two(self):
        """Averaging two different quaternions."""
        quat1 = euler2quat(np.pi / 4, 0, 0)
        quat2 = euler2quat(-np.pi / 4, 0, 0)

        avg_quat = average_quaternions([quat1, quat2])