import pickle
from itertools import product

import pytest

import numpy as np
from scipy.spatial.transform import Rotation, RigidTransform
from scipy.spatial.transform._rigid_transform import normalize_dual_quaternion
from scipy._lib._array_api import (
    is_lazy_array,
    xp_vector_norm,
    is_numpy,
    xp_assert_close,
    make_xp_test_case,
    xp_assert_equal,
    xp_promote
)
import scipy._external.array_api_extra as xpx


def rotation_to_xp(r: Rotation, xp):
    dtype = xpx.default_dtype(xp)
    return Rotation.from_quat(xp.asarray(r.as_quat(), dtype=dtype))


def rigid_transform_to_xp(r: RigidTransform, xp):
    dtype = xpx.default_dtype(xp)
    return RigidTransform.from_matrix(xp.asarray(r.as_matrix(), dtype=dtype))


@make_xp_test_case((RigidTransform, "as_matrix"))
def test_repr(xp):
    actual = repr(RigidTransform.from_matrix(xp.eye(4)))
    expected = """\
RigidTransform.from_matrix(array([[1., 0., 0., 0.],
                                  [0., 1., 0., 0.],
                                  [0., 0., 1., 0.],
                                  [0., 0., 0., 1.]]))"""
    if is_numpy(xp):
        assert actual == expected
    else:
        assert actual.startswith("RigidTransform.from_matrix(")

    tf = RigidTransform.from_matrix(xp.asarray(RigidTransform.identity(2).as_matrix()))
    actual = repr(tf)
    expected = """\
RigidTransform.from_matrix(array([[[1., 0., 0., 0.],
                                   [0., 1., 0., 0.],
                                   [0., 0., 1., 0.],
                                   [0., 0., 0., 1.]],
                           
                                  [[1., 0., 0., 0.],
                                   [0., 1., 0., 0.],
                                   [0., 0., 1., 0.],
                                   [0., 0., 0., 1.]]]))"""
    if is_numpy(xp):
        assert actual == expected
    else:
        assert actual.startswith("RigidTransform.from_matrix(")


@make_xp_test_case((RigidTransform, "from_rotation"))
@pytest.mark.parametrize("ndim", range(1, 4))
def test_from_rotation(xp, ndim: int):
    atol = 1e-12
    rng = np.random.default_rng(0)
    shape = (ndim,) * (ndim - 1) + (4,)
    r = rotation_to_xp(Rotation.from_quat(rng.normal(size=shape)), xp=xp)
    tf = RigidTransform.from_rotation(r)
    xp_assert_close(tf.as_matrix()[..., :3, :3], r.as_matrix(), atol=atol)
    xp_assert_close(tf.as_matrix()[..., :3, 3], xp.zeros(shape[:-1] + (3,)), atol=atol)
    xp_assert_close(tf.as_matrix()[..., 3, :3], xp.zeros(shape[:-1] + (3,)), atol=atol)
    xp_assert_close(tf.as_matrix()[..., 3, 3], xp.ones(shape[:-1]), atol=atol)
    assert tf.single == (ndim == 1)


@make_xp_test_case((RigidTransform, "from_translation"))
@pytest.mark.parametrize("ndim", range(1, 4))
def test_from_translation(xp, ndim: int):
    shape = (ndim,) * (ndim - 1)
    t = xp.reshape(xp.arange(ndim ** (ndim-1) * 3), shape + (3,))
    tf = RigidTransform.from_translation(t)

    expected = xp.tile(xp.eye(4), shape + (1, 1))
    t_float = xp_promote(t, force_floating=True, xp=xp)
    expected = xpx.at(expected)[..., :3, 3].set(t_float)
    xp_assert_close(tf.as_matrix(), expected)
    assert tf.single == (ndim == 1)


def test_from_translation_array_like():
    # Test single translation
    t = [1, 2, 3]
    tf = RigidTransform.from_translation(t)
    tf_expected = RigidTransform.from_translation(np.array(t))
    xp_assert_close(tf.as_matrix(), tf_expected.as_matrix())
    assert tf.single

    # Test multiple translations
    t = [[1, 2, 3], [4, 5, 6]]
    tf = RigidTransform.from_translation(t)
    tf_expected = RigidTransform.from_translation(np.array(t))
    xp_assert_close(tf.as_matrix(), tf_expected.as_matrix())
    assert not tf.single


@make_xp_test_case((RigidTransform, "from_matrix"), (RigidTransform, "as_matrix"))
@pytest.mark.parametrize("ndim", range(1, 4))
def test_from_matrix(xp, ndim: int):
    atol = 1e-12
    shape = (ndim,) * (ndim - 1) + (4, 4)
    dtype = xpx.default_dtype(xp)

    matrix = xp.tile(xp.eye(4), shape[:-2] + (1, 1))
    t = xp.reshape(xp.arange(ndim ** (ndim-1) * 3, dtype=dtype), shape[:-2] + (3,))
    matrix = xpx.at(matrix)[..., :3, 3].set(t)

    tf = RigidTransform.from_matrix(matrix)
    xp_assert_close(tf.as_matrix(), matrix, atol=atol)
    assert tf.single == (ndim == 1)

    # Test non-1 determinant
    matrix = xp.tile(xp.eye(4), shape[:-2] + (1, 1))
    matrix = xpx.at(matrix)[..., :3, :3].set(xp.eye(3) * 2.0)
    tf = RigidTransform.from_matrix(matrix)
    expected = xp.tile(xp.eye(4), shape[:-2] + (1, 1))
    xp_assert_close(tf.as_matrix(), expected, atol=atol)

    # Test non-orthogonal rotation matrix
    matrix = xp.tile(xp.eye(4), shape[:-2] + (1, 1))
    # matrix is equivalent to [[1, 1, 0, 0],
    #                          [0, 1, 0, 0],
    #                          [0, 0, 1, 0],
    #                          [0, 0, 0, 1]]
    matrix = xpx.at(matrix)[..., 0, 1].set(1.0)
    tf = RigidTransform.from_matrix(matrix)
    expected = xp.tile(xp.eye(4), shape[:-2] + (1, 1))
    expected = xpx.at(expected)[..., 0, 0].set(0.894427)
    expected = xpx.at(expected)[..., 0, 1].set(0.447214)
    expected = xpx.at(expected)[..., 1, 0].set(-0.447214)
    expected = xpx.at(expected)[..., 1, 1].set(0.894427)
    xp_assert_close(tf.as_matrix(), expected, atol=1e-6)

    # Test invalid matrix
    invalid = xp.tile(xp.eye(4), shape[:-2] + (1, 1))
    invalid = xpx.at(invalid)[..., 3, 3].set(2)  # Invalid last row
    if is_lazy_array(invalid):
        tf = RigidTransform.from_matrix(invalid)
        assert xp.all(xp.isnan(tf.as_matrix()))
    else:
        with pytest.raises(ValueError):
            RigidTransform.from_matrix(invalid)


def test_from_matrix_array_like():
    # Test single transform matrix
    matrix = [[1, 0, 0, 0],
              [0, 1, 0, 0],
              [0, 0, 1, 0],
              [0, 0, 0, 1]]
    expected = np.eye(4)
    tf = RigidTransform.from_matrix(matrix)
    xp_assert_close(tf.as_matrix(), expected)
    assert tf.single

    # Test multiple transform matrices
    matrices = [matrix, matrix]
    tf = RigidTransform.from_matrix(matrices)
    for i in range(len(matrices)):
        xp_assert_close(tf.as_matrix()[i, ...], expected)
    assert not tf.single


@make_xp_test_case((RigidTransform, "from_components"))
@pytest.mark.parametrize("r_ndim", range(1, 3))
@pytest.mark.parametrize("t_ndim", range(1, 3))
def test_from_components(xp, r_ndim: int, t_ndim: int):
    atol = 1e-12
    dims = (6, 5, 4, 3)  # Common shape
    q_shape = dims[:r_ndim - 1][::-1] + (4,)
    t_shape = dims[:t_ndim - 1][::-1] + (3,)
    tf_shape = np.broadcast_shapes(q_shape[:-1], t_shape[:-1]) + (4, 4)
    rng = np.random.default_rng(0)

    t = xp.reshape(xp.arange(np.prod(t_shape[:-1]) * 3), t_shape)
    r = rotation_to_xp(Rotation.from_quat(rng.random(size=q_shape)), xp=xp)
    tf = RigidTransform.from_components(t, r)

    expected = xp.zeros(tf_shape)
    expected = xpx.at(expected)[..., :3, :3].set(r.as_matrix())
    t_float = xp_promote(t, force_floating=True, xp=xp)
    expected = xpx.at(expected)[..., :3, 3].set(t_float)
    expected = xpx.at(expected)[..., 3, 3].set(1)
    xp_assert_close(tf.as_matrix(), expected, atol=atol)
    assert tf.single == (r_ndim == 1 and t_ndim == 1)


def test_from_components_array_like():
    rng = np.random.default_rng(123)
    # Test single rotation and translation
    t = [1, 2, 3]
    r = Rotation.random(rng=rng)
    tf = RigidTransform.from_components(t, r)
    tf_expected = RigidTransform.from_components(np.array(t), r)
    xp_assert_close(tf.as_matrix(), tf_expected.as_matrix(), atol=1e-12)
    assert tf.single

    # Test multiple rotations and translations
    t = [[1, 2, 3], [4, 5, 6]]
    r = Rotation.random(len(t), rng=rng)
    tf = RigidTransform.from_components(t, r)
    tf_expected = RigidTransform.from_components(np.array(t), r)
    xp_assert_close(tf.as_matrix(), tf_expected.as_matrix(), atol=1e-12)
    assert not tf.single


@make_xp_test_case((RigidTransform, "as_components"))
@pytest.mark.parametrize("ndim", range(1, 4))
def test_as_components(xp, ndim):
    dtype = xpx.default_dtype(xp)
    atol = 1e-12 if dtype == xp.float64 else 1e-6
    shape = (ndim,) * (ndim - 1)
    rng = np.random.default_rng(123)
    t = xp.asarray(rng.normal(size=shape + (3,)), dtype=dtype)
    r = rotation_to_xp(Rotation.from_quat(rng.random(shape + (4,))), xp=xp)
    tf = RigidTransform.from_components(t, r)
    new_t, new_r = tf.as_components()
    assert xp.all(new_r.approx_equal(r, atol=atol))
    xp_assert_close(new_t, t, atol=atol)


@make_xp_test_case((RigidTransform, "from_exp_coords"))
@pytest.mark.parametrize("dim", range(1, 4))
def test_from_exp_coords(xp, dim: int):
    shape = (dim,) * (dim - 1)
    # example from 3.3 of
    # https://hades.mech.northwestern.edu/images/2/25/MR-v2.pdf
    dtype = xpx.default_dtype(xp)
    angle1 = np.deg2rad(30.0)
    mat = xp.asarray([
        [np.cos(angle1), -np.sin(angle1), 0.0, 1.0],
        [np.sin(angle1), np.cos(angle1), 0.0, 2.0],
        [0.0, 0.0, 1.0, 0.0],
        [0.0, 0.0, 0.0, 1.0]
    ], dtype=dtype)
    mat = xp.tile(mat, shape + (1, 1))
    tf1 = RigidTransform.from_matrix(mat)
    angle2 = np.deg2rad(60.0)
    mat = xp.asarray([
        [np.cos(angle2), -np.sin(angle2), 0.0, 2.0],
        [np.sin(angle2), np.cos(angle2), 0.0, 1.0],
        [0.0, 0.0, 1.0, 0.0],
        [0.0, 0.0, 0.0, 1.0]
    ], dtype=dtype)
    mat = xp.tile(mat, shape + (1, 1))
    tf2 = RigidTransform.from_matrix(mat)
    expected = tf2 * tf1.inv()
    deg2rag = xp.asarray(np.deg2rad(30.0))
    exp_coords = deg2rag * xp.asarray([0.0, 0.0, 1.0, 3.37, -3.37, 0.0])
    exp_coords = xp.tile(exp_coords, shape + (1,))
    actual = RigidTransform.from_exp_coords(exp_coords)
    xp_assert_close(actual.as_matrix(), expected.as_matrix(), atol=1e-2)

    # test cases generated by comparison to pytransform3d
    exp_coords = xp.asarray([
        [-2.01041204, -0.52983629, 0.65773501,
         0.10386614, 0.05855009, 0.54959179],
        [-0.22537438, -0.24132627, -2.4747121,
         -0.09158594,  1.88075832, -0.03197204]
    ])
    exp_coords = xp.tile(exp_coords, shape + (1,))
    expected_matrix = xp.asarray([
        [[0.76406621, 0.10504613, -0.63652819, -0.10209961],
         [0.59956454, -0.47987325, 0.64050295, 0.40158789],
         [-0.2381705, -0.87102639, -0.42963687, 0.19637636],
         [0., 0., 0., 1.]],
        [[-0.78446989, 0.61157488, 0.10287448, 1.33330055],
         [-0.58017785, -0.78232107, 0.22664378, 0.52660831],
         [0.21909052, 0.11810973, 0.96852952, -0.02968529],
         [0., 0., 0., 1.]]
    ])
    expected_matrix = xp.tile(expected_matrix, shape + (1, 1))
    xp_assert_close(
        RigidTransform.from_exp_coords(exp_coords).as_matrix(),
        expected_matrix, atol=1e-8)

    # identity
    expected_matrix = xp.tile(xp.eye(4), shape + (1, 1))
    exp_coords = xp.zeros(shape + (6,), dtype=dtype)
    xp_assert_close(
        RigidTransform.from_exp_coords(exp_coords).as_matrix(),
        expected_matrix, atol=1e-12)

    # only translation
    expected_matrix = xp.asarray([
        [[1.0, 0.0, 0.0, 3.0],
         [0.0, 1.0, 0.0, -5.4],
         [0.0, 0.0, 1.0, 100.2],
         [0.0, 0.0, 0.0, 1.0]],
        [[1.0, 0.0, 0.0, -3.0],
         [0.0, 1.0, 0.0, 13.3],
         [0.0, 0.0, 1.0, 1.3],
         [0.0, 0.0, 0.0, 1.0]]
    ])
    expected_matrix = xp.tile(expected_matrix, shape + (1, 1, 1))
    exp_coords = xp.asarray([
        [0.0, 0.0, 0.0, 3.0, -5.4, 100.2],
        [0.0, 0.0, 0.0, -3.0, 13.3, 1.3],
    ])
    exp_coords = xp.tile(exp_coords, shape + (1, 1))
    actual = RigidTransform.from_exp_coords(exp_coords)
    xp_assert_close(actual.as_matrix(), expected_matrix, atol=1e-12)

    # only rotation
    angles = xp.asarray([[34, -12, 0.5], [-102, -55, 30]])
    angles = xp.tile(angles, shape + (1, 1))
    rot = Rotation.from_euler('zyx', angles, degrees=True)
    rotvec = rot.as_rotvec()
    expected_matrix = xp.tile(xp.eye(4), shape + (2, 1, 1))
    expected_matrix = xpx.at(expected_matrix)[..., :3, :3].set(rot.as_matrix())
    exp_coords = xp.concat((rotvec, xp.zeros_like(rotvec)), axis=-1)
    actual = RigidTransform.from_exp_coords(exp_coords)
    xp_assert_close(actual.as_matrix(), expected_matrix, atol=1e-12)


def test_from_exp_coords_array_like():
    rng = np.random.default_rng(123)
    # Test single transform
    t = np.array([1, 2, 3])
    r = Rotation.random(rng=rng)
    tf_expected = RigidTransform.from_components(t, r)
    exp_coords = tf_expected.as_exp_coords().tolist()
    assert isinstance(exp_coords, list)
    tf = RigidTransform.from_exp_coords(exp_coords)
    xp_assert_close(tf.as_matrix(), tf_expected.as_matrix(), atol=1e-12)

    # Test multiple transforms
    t = [[1, 2, 3], [4, 5, 6]]
    r = Rotation.random(len(t), rng=rng)
    tf_expected = RigidTransform.from_components(t, r)
    exp_coords = tf_expected.as_exp_coords().tolist()
    assert isinstance(exp_coords, list)
    tf = RigidTransform.from_exp_coords(exp_coords)
    xp_assert_close(tf.as_matrix(), tf_expected.as_matrix(), atol=1e-12)


@make_xp_test_case((RigidTransform, "as_exp_coords"))
@pytest.mark.parametrize("ndim", range(1, 4))
def test_as_exp_coords(xp, ndim: int):
    shape = (ndim,) * (ndim - 1)
    # identity
    expected = xp.zeros(shape + (6,))
    actual = RigidTransform.from_exp_coords(expected).as_exp_coords()
    xp_assert_close(actual, expected, atol=1e-12)

    rng = np.random.default_rng(10)

    # pure rotation
    rot_vec = xp.asarray(rng.normal(scale=0.1, size=shape + (1000, 3)))
    tf = RigidTransform.from_rotation(Rotation.from_rotvec(rot_vec))
    exp_coords = tf.as_exp_coords()
    xp_assert_close(exp_coords[..., :3], rot_vec, atol=1e-12)
    expected = xp.zeros_like(rot_vec)
    xp_assert_close(exp_coords[..., 3:], expected, atol=1e-16)

    # pure translation
    translation = xp.asarray(rng.normal(scale=100.0, size=shape + (1000, 3)))
    tf = RigidTransform.from_translation(translation)
    exp_coords = tf.as_exp_coords()
    xp_assert_close(exp_coords[..., :3], expected, atol=1e-16)
    xp_assert_close(exp_coords[..., 3:], translation, atol=1e-15)


@make_xp_test_case((RigidTransform, "from_dual_quat"))
@pytest.mark.parametrize("ndim", range(1, 4))
def test_from_dual_quat(xp, ndim: int):
    dtype = xpx.default_dtype(xp)
    atol = 1e-12 if dtype == xp.float64 else 1e-7
    shape = (ndim,) * (ndim - 1)

    # identity
    dq = xp.asarray([0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0], dtype=dtype)
    dq = xp.tile(dq, shape + (1,))
    expected = xp.tile(xp.eye(4), shape + (1, 1))
    xp_assert_close(RigidTransform.from_dual_quat(dq).as_matrix(), expected, atol=atol)
    dq = xp.asarray([1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], dtype=dtype)
    dq = xp.tile(dq, shape + (1,))
    xp_assert_close(RigidTransform.from_dual_quat(dq, scalar_first=True).as_matrix(),
                    expected, atol=atol)

    # only translation
    dq = xp.asarray([0, 0, 0, 1, 0.25, 0.15, -0.7, 0], dtype=dtype)
    dq = xp.tile(dq, shape + (1,))
    actual = RigidTransform.from_dual_quat(dq)
    expected_matrix = xp.asarray([
        [1, 0, 0, 0.5],
        [0, 1, 0, 0.3],
        [0, 0, 1, -1.4],
        [0, 0, 0, 1]
    ])
    expected_matrix = xp.tile(expected_matrix, shape + (1, 1))
    xp_assert_close(actual.as_matrix(), expected_matrix, atol=atol)
    dq = xp.asarray([1, 0, 0, 0, 0, 0.25, 0.15, -0.7], dtype=dtype)
    dq = xp.tile(dq, shape + (1,))
    actual = RigidTransform.from_dual_quat(dq, scalar_first=True)
    xp_assert_close(actual.as_matrix(), expected_matrix, atol=atol)

    # only rotation
    angles = xp.asarray([65, -13, 90], dtype=dtype)
    angles = xp.tile(angles, shape + (1,))
    actual_rot = Rotation.from_euler("xyz", angles, degrees=True)
    qrot = actual_rot.as_quat()
    dq = xp.concat((qrot, xp.zeros_like(qrot)), axis=-1)
    actual = RigidTransform.from_dual_quat(dq)
    expected_matrix = xp.tile(xp.eye(4), shape + (1, 1))
    expected_matrix = xpx.at(expected_matrix)[..., :3, :3].set(actual_rot.as_matrix())
    xp_assert_close(actual.as_matrix(), expected_matrix, atol=atol)

    qrot = actual_rot.as_quat(scalar_first=True)
    dq = xp.concat((qrot, xp.zeros_like(qrot)), axis=-1)
    actual = RigidTransform.from_dual_quat(dq, scalar_first=True)
    expected_matrix = xp.tile(xp.eye(4), shape + (1, 1))
    expected_matrix = xpx.at(expected_matrix)[..., :3, :3].set(actual_rot.as_matrix())
    xp_assert_close(actual.as_matrix(), expected_matrix, atol=atol)

    # rotation and translation
    # rtol is set to 1e-7 because xp_assert_close deviates from
    # np.testing.assert_allclose in that it does not automatically default to 1e-7 for
    # floating point inputs.
    # See https://numpy.org/doc/2.2/reference/generated/numpy.testing.assert_allclose.html
    dq = xp.asarray(
        [[0.0617101, -0.06483886, 0.31432811, 0.94508498,
          0.04985168, -0.26119618, 0.1691491, -0.07743254],
         [0.19507259, 0.49404931, -0.06091285, 0.8450749,
          0.65049656, -0.30782513, 0.16566752, 0.04174544]])
    dq = xp.tile(dq, shape + (1, 1))
    actual = RigidTransform.from_dual_quat(dq)
    expected_matrix = xp.asarray(
        [[[0.79398752, -0.60213598, -0.08376202, 0.24605262],
          [0.58613113, 0.79477941, -0.15740392, -0.4932833],
          [0.16135089, 0.07588122, 0.98397557, 0.34262676],
          [0., 0., 0., 1.]],
         [[0.50440981, 0.2957028, 0.81125249, 1.20934468],
          [0.08979911, 0.91647262, -0.3898898, -0.70540077],
          [-0.8587822, 0.26951399, 0.43572393, -0.47776265],
          [0., 0., 0., 1.]]])
    expected_matrix = xp.tile(expected_matrix, shape + (1, 1, 1))
    xp_assert_close(actual.as_matrix(), expected_matrix, atol=atol, rtol=1e-7)

    dq = xp.asarray(
        [[0.94508498, 0.0617101, -0.06483886, 0.31432811,
          -0.07743254, 0.04985168, -0.26119618, 0.1691491],
         [0.8450749, 0.19507259, 0.49404931, -0.06091285,
          0.04174544, 0.65049656, -0.30782513, 0.16566752]])
    dq = xp.tile(dq, shape + (1, 1))
    actual = RigidTransform.from_dual_quat(dq, scalar_first=True)
    xp_assert_close(actual.as_matrix(), expected_matrix, atol=atol, rtol=1e-7)

    # unnormalized dual quaternions

    # invalid real quaternion with norm 0
    dq = xp.zeros(shape + (8,))
    actual = RigidTransform.from_dual_quat(dq)
    expected = xp.tile(xp.eye(4), shape + (1, 1))
    xp_assert_close(actual.as_matrix(), expected, atol=atol)

    # real quaternion with norm != 1
    unnormalized_dual_quat = xp.asarray(
        [-0.2547655, 1.23506123, 0.20230088, 0.24247194,  # norm 1.3
         0.38559628, 0.08184063, 0.1755943, -0.1582222]  # orthogonal
    )
    xp_assert_close(xp_vector_norm(unnormalized_dual_quat[:4]), xp.asarray(1.3)[()],
                    atol=atol)
    xp_assert_close(xp.vecdot(unnormalized_dual_quat[:4],
                              unnormalized_dual_quat[4:])[()],
                    xp.asarray(0.0)[()], atol=1e-8)

    dq = xp.tile(unnormalized_dual_quat, shape + (1,))
    dual_quat = RigidTransform.from_dual_quat(dq).as_dual_quat()

    expected_ones = xp.ones(shape) if shape != () else xp.asarray(1.0)[()]
    expected_zeros = xp.zeros(shape) if shape != () else xp.asarray(0.0)[()]
    xp_assert_close(xp_vector_norm(dual_quat[..., :4], axis=-1), expected_ones,
                    atol=1e-12)
    vecdot = xp.vecdot(dual_quat[..., :4], dual_quat[..., 4:])
    vecdot = vecdot[()] if vecdot.shape == () else vecdot
    xp_assert_close(vecdot, expected_zeros, atol=atol)

    # real and dual quaternion are not orthogonal
    unnormalized_dual_quat = xp.asarray(
        [0.20824458, 0.75098079, 0.54542913, -0.30849493,  # unit norm
         -0.16051025, 0.10742978, 0.21277201, 0.20596935]  # not orthogonal
    )
    xp_assert_close(xp_vector_norm(unnormalized_dual_quat[:4]), xp.asarray(1.0)[()],
                    atol=atol)
    assert xp.vecdot(unnormalized_dual_quat[:4], unnormalized_dual_quat[4:]) != 0.0
    dq = xp.tile(unnormalized_dual_quat, shape + (1,))
    dual_quat = RigidTransform.from_dual_quat(dq).as_dual_quat()

    xp_assert_close(xp_vector_norm(dual_quat[..., :4], axis=-1), expected_ones,
                    atol=1e-12)
    vecdot = xp.vecdot(dual_quat[..., :4], dual_quat[..., 4:])
    vecdot = vecdot[()] if vecdot.shape == () else vecdot
    xp_assert_close(vecdot, expected_zeros, atol=atol)

    # invalid real quaternion with norm 0, non-orthogonal dual quaternion
    unnormalized_dual_quat = xp.asarray(
        [0.0, 0.0, 0.0, 0.0, -0.16051025, 0.10742978, 0.21277201, 0.20596935])
    assert xp.vecdot(xp.asarray([0.0, 0, 0, 1]), unnormalized_dual_quat[4:]) != 0.0
    dq = xp.tile(unnormalized_dual_quat, shape + (1,))
    dual_quat = RigidTransform.from_dual_quat(dq).as_dual_quat()

    xp_assert_close(xp_vector_norm(dual_quat[..., :4], axis=-1), expected_ones,
                    atol=1e-12)
    vecdot = xp.vecdot(dual_quat[..., :4], dual_quat[..., 4:])
    vecdot = vecdot[()] if vecdot.shape == () else vecdot
    xp_assert_close(vecdot, expected_zeros, atol=atol)

    # compensation for precision loss in real quaternion
    rng = np.random.default_rng(1000)
    t = xp.asarray(rng.normal(size=shape + (3,)), dtype=dtype)
    q = xp.asarray(rng.normal(size=shape + (4,)), dtype=dtype)
    r = Rotation.from_quat(q)
    random_dual_quats = RigidTransform.from_components(t, r).as_dual_quat()

    # ensure that random quaternions are not normalized
    random_dual_quats = xpx.at(random_dual_quats)[..., :4].add(0.01)
    assert not xp.any(xpx.isclose(xp_vector_norm(random_dual_quats[..., :4], axis=-1),
                                  1.0, atol=0.0001))
    dual_quat_norm = RigidTransform.from_dual_quat(
        random_dual_quats).as_dual_quat()
    xp_assert_close(xp_vector_norm(dual_quat_norm[..., :4], axis=-1), expected_ones,
                    atol=atol)

    # compensation for precision loss in dual quaternion, results in violation
    # of orthogonality constraint
    t = xp.asarray(rng.normal(size=shape + (3,)), dtype=dtype)
    q = xp.asarray(rng.normal(size=shape + (4,)), dtype=dtype)
    r = Rotation.from_quat(q)
    random_dual_quats = RigidTransform.from_components(t, r).as_dual_quat()

    # ensure that random quaternions are not normalized
    random_dual_quats = xpx.at(random_dual_quats)[..., 4:].add(0.1)
    q_norm = xp.vecdot(random_dual_quats[..., :4], random_dual_quats[..., 4:])
    assert not xp.any(xpx.isclose(q_norm, 0.0, atol=0.0001))
    dual_quat_norm = RigidTransform.from_dual_quat(
        random_dual_quats).as_dual_quat()
    vecdot = xp.vecdot(dual_quat[..., :4], dual_quat[..., 4:])
    vecdot = vecdot[()] if vecdot.shape == () else vecdot
    xp_assert_close(vecdot, expected_zeros, atol=atol)
    xp_assert_close(random_dual_quats[..., :4], dual_quat_norm[..., :4], atol=atol)


def test_from_dual_quat_array_like():
    rng = np.random.default_rng(123)
    # Test single transform
    t = np.array([1, 2, 3])
    r = Rotation.random(rng=rng)
    tf_expected = RigidTransform.from_components(t, r)
    dual_quat = tf_expected.as_dual_quat().tolist()
    assert isinstance(dual_quat, list)
    tf = RigidTransform.from_dual_quat(dual_quat)
    xp_assert_close(tf.as_matrix(), tf_expected.as_matrix(), atol=1e-12)

    # Test multiple transforms
    t = [[1, 2, 3], [4, 5, 6]]
    r = Rotation.random(len(t), rng=rng)
    tf_expected = RigidTransform.from_components(t, r)
    dual_quat = tf_expected.as_dual_quat().tolist()
    assert isinstance(dual_quat, list)
    tf = RigidTransform.from_dual_quat(dual_quat)
    xp_assert_close(tf.as_matrix(), tf_expected.as_matrix(), atol=1e-12)


@make_xp_test_case((RigidTransform, "as_dual_quat"))
@pytest.mark.parametrize("ndim", range(1, 4))
def test_as_dual_quat(xp, ndim: int):
    dtype = xpx.default_dtype(xp)
    shape = (ndim,) * (ndim - 1)
    # identity
    expected = xp.asarray([0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0], dtype=dtype)
    actual = rigid_transform_to_xp(RigidTransform.identity(), xp).as_dual_quat()
    xp_assert_close(actual, expected, atol=1e-12)

    expected = xp.asarray([1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0])
    tf = rigid_transform_to_xp(RigidTransform.identity(), xp)
    actual = tf.as_dual_quat(scalar_first=True)
    xp_assert_close(actual, expected, atol=1e-12)

    rng = np.random.default_rng(10)

    # only rotation
    for _ in range(10):
        q = xp.asarray(rng.normal(size=shape + (4,)), dtype=dtype)
        real_part = Rotation.from_quat(q).as_quat()
        dual_part = xp.zeros_like(real_part)
        expected = xp.concat((real_part, dual_part), axis=-1)
        actual = RigidTransform.from_dual_quat(expected).as_dual_quat()
        # because of double cover:
        actual = actual * xp.sign(actual[..., 0, None])
        expected = expected * xp.sign(expected[..., 0, None])
        xp_assert_close(actual, expected, atol=1e-12)

    # only translation
    for _ in range(10):
        tf = 0.5 * xp.asarray(rng.normal(size=shape + (3,)), dtype=dtype)
        expected = xp.zeros(shape + (8,), dtype=dtype)
        expected = xpx.at(expected)[..., 3].set(1.0)
        expected = xpx.at(expected)[..., 4:7].set(tf)
        actual = RigidTransform.from_dual_quat(expected).as_dual_quat()
        # because of double cover:
        actual = actual * xp.sign(actual[..., 0, None])
        expected = expected * xp.sign(expected[..., 0, None])
        xp_assert_close(actual, expected, atol=1e-12)

    # rotation and translation
    for _ in range(10):
        t = xp.asarray(rng.normal(size=shape + (3,)), dtype=dtype)
        r = Rotation.from_quat(xp.asarray(rng.normal(size=shape + (4,)), dtype=dtype))
        expected = RigidTransform.from_components(t, r).as_dual_quat()
        actual = RigidTransform.from_dual_quat(expected).as_dual_quat()
        # because of double cover:
        actual = actual * xp.sign(actual[..., 0, None])
        expected = expected * xp.sign(expected[..., 0, None])
        xp_assert_close(actual, expected, atol=1e-12)


@make_xp_test_case(
    (RigidTransform, "from_components"), (RigidTransform, "as_components"),
    (RigidTransform, "from_exp_coords"), (RigidTransform, "as_exp_coords"),
    (RigidTransform, "from_matrix"), (RigidTransform, "as_matrix"),
    (RigidTransform, "from_dual_quat"), (RigidTransform, "as_dual_quat")
)
@pytest.mark.parametrize("ndim", range(1, 4))
def test_from_as_internal_consistency(xp, ndim: int):
    dtype = xpx.default_dtype(xp)
    atol = 1e-12
    n = 10
    rng = np.random.default_rng(10)
    shape = (n,) + (ndim,) * (ndim - 1)

    t = xp.asarray(rng.normal(size=shape + (3,)), dtype=dtype)
    r = Rotation.from_quat(xp.asarray(rng.normal(size=shape + (4,)), dtype=dtype))
    tf0 = RigidTransform.from_components(t, r)

    tf1 = RigidTransform.from_components(*tf0.as_components())
    xp_assert_close(tf0.as_matrix(), tf1.as_matrix(), atol=atol)

    tf1 = RigidTransform.from_components(tf0.translation, tf0.rotation)
    xp_assert_close(tf0.as_matrix(), tf1.as_matrix(), atol=atol)

    tf1 = RigidTransform.from_exp_coords(tf0.as_exp_coords())
    xp_assert_close(tf0.as_matrix(), tf1.as_matrix(), atol=atol)

    tf1 = RigidTransform.from_matrix(tf0.as_matrix())
    xp_assert_close(tf0.as_matrix(), tf1.as_matrix(), atol=atol)

    tf1 = RigidTransform.from_dual_quat(tf0.as_dual_quat())
    xp_assert_close(tf0.as_matrix(), tf1.as_matrix(), atol=atol)

    # exp_coords small rotation
    t = xp.asarray(rng.normal(scale=1000.0, size=shape + (3,)), dtype=dtype)
    rotvec = xp.asarray(rng.normal(scale=1e-10, size=shape + (3,)), dtype=dtype)
    r = Rotation.from_rotvec(rotvec)
    tf0 = RigidTransform.from_components(t, r)
    tf1 = RigidTransform.from_exp_coords(tf0.as_exp_coords())
    xp_assert_close(tf0.as_matrix(), tf1.as_matrix(), atol=atol)


def test_identity():
    # We do not use xp here because identity always returns numpy arrays
    atol = 1e-12
    # Test single identity
    tf = RigidTransform.identity()
    xp_assert_close(tf.as_matrix(), np.eye(4), atol=atol)
    # Test multiple identities
    tf = RigidTransform.identity(5)
    xp_assert_close(tf.as_matrix(), np.array([np.eye(4)] * 5), atol=atol)
    # Test shape
    tf = RigidTransform.identity(shape=3)
    expected = np.tile(np.eye(4), (3, 1, 1))
    xp_assert_close(tf.as_matrix(), expected, atol=atol)
    tf = RigidTransform.identity(shape=(2, 3))
    expected = np.tile(np.eye(4), (2, 3, 1, 1))
    xp_assert_close(tf.as_matrix(), expected, atol=atol)
    # Test errors
    with pytest.raises(ValueError, match="Only one of `num` and `shape` can be."):
        RigidTransform.identity(10, shape=(2, 3))
    with pytest.raises(TypeError, match="takes from 0 to 1 positional arguments"):
        RigidTransform.identity(None, (-1, 3))  # Shape is kwarg only
    with pytest.raises(ValueError, match="`shape` must be an int or a tuple of ints"):
        RigidTransform.identity(shape="invalid")


@make_xp_test_case((RigidTransform, "apply"))
def test_apply(xp):
    atol = 1e-12
    # Broadcast shape: (6, 5, 4, 2) ( + (3,) for vectors, + (4,) for rotations)
    vector_shapes = [(), (1,), (2,), (1, 2), (5, 1, 2)]
    tf_shapes = [(), (1,), (2,), (1, 2), (4, 2), (1, 4, 2), (5, 4, 2), (6, 5, 4, 2)]
    rng = np.random.default_rng(123)

    for tf_shape, v_shape in product(tf_shapes, vector_shapes):
        # Random rotation and translation
        t = xp.asarray(rng.normal(size=tf_shape + (3,)))
        q = xp.asarray(rng.normal(size=tf_shape + (4,)))
        r = Rotation.from_quat(q)
        tf = RigidTransform.from_components(t, r)

        vecs = xp.asarray(rng.normal(size=v_shape + (3,)))
        expected = t + r.apply(vecs)
        res = tf.apply(vecs)
        assert res.shape == np.broadcast_shapes(tf_shape, v_shape) + (3,)
        xp_assert_close(res, expected, atol=atol)


def test_apply_array_like():
    rng = np.random.default_rng(123)
    # Single vector
    t = np.array([1, 2, 3])
    r = Rotation.random(rng=rng)
    tf = RigidTransform.from_components(t, r)
    vec = [1, 0, 0]
    expected = t + r.apply(vec)
    xp_assert_close(tf.apply(vec), expected, atol=1e-12)

    # Multiple vectors
    t = np.array([[1, 2, 3], [4, 5, 6]])
    r = Rotation.random(len(t), rng=rng)
    tf = RigidTransform.from_components(t, r)
    vec = [[1, 0, 0], [0, 1, 0]]
    expected = t + r.apply(vec)
    xp_assert_close(tf.apply(vec), expected, atol=1e-12)


def test_apply_matrix_equivalence():
    """Test documented equivalence for single transform:
    `apply(vector) == translation + vector @ rotation.as_matrix().T.`"""
    t = np.array([1.0, 2.0, 3.0])
    r = Rotation.from_rotvec([0, 0, 1])
    tf = RigidTransform.from_components(t, r)
    # Single vector (3,)
    v = np.array([1.0, 0.0, 0.0])
    xp_assert_close(tf.apply(v), t + v @ r.as_matrix().T)
    # Multiple vectors (P, 3)
    arr = np.array([[1, 0, 0], [0, 1, 0]], dtype=float)
    xp_assert_close(tf.apply(arr), t + arr @ r.as_matrix().T)


@make_xp_test_case((RigidTransform, "apply"))
def test_inverse_apply(xp):
    atol = 1e-12
    # Broadcast shape: (6, 5, 4, 2) ( + (3,) for vectors, + (4,) for rotations)
    vector_shapes = [(), (1,), (2,), (1, 2), (5, 1, 2)]
    tf_shapes = [(), (1,), (2,), (1, 2), (4, 2), (1, 4, 2), (5, 4, 2), (6, 5, 4, 2)]
    rng = np.random.default_rng(123)

    for tf_shape, v_shape in product(tf_shapes, vector_shapes):
        # Random rotation and translation
        t = xp.asarray(rng.normal(size=tf_shape + (3,)))
        q = xp.asarray(rng.normal(size=tf_shape + (4,)))
        r = Rotation.from_quat(q)
        tf = RigidTransform.from_components(t, r)

        vecs = xp.asarray(rng.normal(size=v_shape + (3,)))
        expected = tf.inv().apply(vecs)
        res = tf.apply(vecs, inverse=True)
        assert res.shape == np.broadcast_shapes(tf_shape, v_shape) + (3,)
        xp_assert_close(res, expected, atol=atol)


@make_xp_test_case((RigidTransform, "apply"))
def test_rotation_alone(xp):
    atol = 1e-12

    r = Rotation.from_euler('z', xp.asarray(90), degrees=True)
    tf = RigidTransform.from_rotation(r)
    vec = xp.asarray([1, 0, 0])
    expected = r.apply(vec)
    xp_assert_close(tf.apply(vec), expected, atol=atol)


@make_xp_test_case((RigidTransform, "apply"))
def test_translation_alone(xp):
    atol = 1e-12
    t = xp.asarray([1.0, 2, 3])
    tf = RigidTransform.from_translation(t)
    vec = xp.asarray([5.0, 6, 7])
    expected = t + vec
    xp_assert_close(tf.apply(vec), expected, atol=atol)


@make_xp_test_case((RigidTransform, "apply"), (RigidTransform, "__mul__"))
def test_composition(xp):
    atol = 1e-12
    tf_shapes = [(), (1,), (2,), (1, 2), (4, 2), (5, 4, 2)]
    dtype = xpx.default_dtype(xp)
    rng = np.random.default_rng(123)

    for tf_shape1, tf_shape2 in product(tf_shapes, repeat=2):
        # Random rotation and translation
        t1 = xp.asarray(rng.normal(size=tf_shape1 + (3,)))
        q1 = xp.asarray(rng.normal(size=tf_shape1 + (4,)))
        r1 = Rotation.from_quat(q1)
        tf1 = RigidTransform.from_components(t1, r1)

        t2 = xp.asarray(rng.normal(size=tf_shape2 + (3,)))
        q2 = xp.asarray(rng.normal(size=tf_shape2 + (4,)))
        r2 = Rotation.from_quat(q2)
        tf2 = RigidTransform.from_components(t2, r2)

        composed = tf2 * tf1
        vec = xp.asarray(rng.normal(size=(3,)), dtype=dtype)
        expected = tf2.apply(tf1.apply(vec))
        res = composed.apply(vec)
        assert res.shape == np.broadcast_shapes(tf_shape1, tf_shape2) + (3,)
        xp_assert_close(res, expected, atol=atol)

        expected = t2 + r2.apply(t1 + r1.apply(vec))
        xp_assert_close(composed.apply(vec), expected, atol=atol)
        assert composed.single == (tf1.single and tf2.single)


@make_xp_test_case((RigidTransform, "__pow__"), (RigidTransform, "__mul__"))
@pytest.mark.parametrize("ndim", range(1, 4))
def test_pow(xp, ndim: int):
    dtype = xpx.default_dtype(xp)
    atol = 1e-12 if dtype == xp.float64 else 1e-6
    num = 10
    rng = np.random.default_rng(100)
    shape = (num,) + (ndim,) * (ndim - 1)
    t = xp.asarray(rng.normal(size=shape + (3,)), dtype=dtype)
    q = xp.asarray(rng.normal(size=shape + (4,)), dtype=dtype)
    r = Rotation.from_quat(q)
    p = RigidTransform.from_components(t, r)
    p_inv = p.inv()

    # Test the short-cuts and other integers
    for n in [-5, -2, -1, 0, 1, 2, 5]:
        q = p**n
        # Regression test for gh-24436
        assert isinstance(q._matrix, type(p._matrix))
        r = RigidTransform.from_matrix(xp.tile(xp.eye(4), shape + (1, 1)))
        for _ in range(abs(n)):
            if n > 0:
                r = r * p
            else:
                r = r * p_inv
        xp_assert_close(q.as_matrix(), r.as_matrix(), atol=atol)

    # Test shape preservation of single
    single_tf = RigidTransform.identity()
    assert (single_tf**n).as_matrix().shape == (4, 4)

    # Test fractional powers
    q = p**0.5
    xp_assert_close((q * q).as_matrix(), p.as_matrix(), atol=atol)
    q = p**-0.5
    xp_assert_close((q * q).as_matrix(), p.inv().as_matrix(), atol=atol)
    q = p** 1.5
    xp_assert_close((q * q).as_matrix(), (p**3).as_matrix(), atol=atol)
    q = p** -1.5
    xp_assert_close((q * q).as_matrix(), (p**-3).as_matrix(), atol=atol)

    # pow function
    tf = pow(RigidTransform.from_matrix(xp.eye(4)), 2)
    xp_assert_close(tf.as_matrix(), xp.eye(4), atol=atol)


@make_xp_test_case((RigidTransform, "__pow__"))
@pytest.mark.parametrize("ndim", range(1, 4))
def test_pow_equivalence_with_rotation(xp, ndim: int):
    atol = 1e-12
    num = 10
    rng = np.random.default_rng(100)
    dtype = xpx.default_dtype(xp)
    shape = (num,) + (ndim,) * (ndim - 1)

    r = Rotation.from_quat(xp.asarray(rng.normal(size=shape + (4,)), dtype=dtype))
    p = RigidTransform.from_rotation(r)
    for n in [-5, -2, -1.5, -1, -0.5, 0.0, 0.5, 1, 1.5, 2, 5]:
        xp_assert_close((p**n).rotation.as_matrix(), (r**n).as_matrix(), atol=atol)


@make_xp_test_case((RigidTransform, "inv"), (RigidTransform, "__mul__"))
@pytest.mark.parametrize("ndim", range(1, 4))
def test_inverse(xp, ndim: int):
    dtype = xpx.default_dtype(xp)
    atol = 1e-12 if dtype == xp.float64 else 1e-6
    rng = np.random.default_rng(100)
    shape = (ndim,) * (ndim - 1)

    # Test inverse transform
    r = Rotation.from_quat(xp.asarray(rng.normal(size=shape + (4,)), dtype=dtype))
    t = xp.asarray(rng.normal(size=shape + (3,)), dtype=dtype)
    tf = RigidTransform.from_components(t, r)

    # Test that tf * tf.inv() equals identity
    tf_inv = tf.inv()
    composed = tf * tf_inv
    expected = xp.tile(xp.eye(4), shape + (1, 1))
    xp_assert_close(composed.as_matrix(), expected, atol=atol)


@make_xp_test_case((RigidTransform, "as_matrix"))
@pytest.mark.parametrize("ndim", range(1, 4))
def test_properties(xp, ndim: int):
    atol = 1e-12 if xpx.default_dtype(xp) == xp.float64 else 1e-6
    shape = (ndim,) * (ndim - 1)
    dtype = xpx.default_dtype(xp)
    rng = np.random.default_rng(100)

    # Test rotation and translation properties
    r = Rotation.from_quat(xp.asarray(rng.normal(size=shape + (4,)), dtype=dtype))
    t = xp.asarray(rng.normal(size=shape + (3,)), dtype=dtype)
    tf = RigidTransform.from_components(t, r)

    xp_assert_close(tf.rotation.as_matrix(), r.as_matrix(), atol=atol)
    assert xp.all(tf.rotation.approx_equal(r, atol=atol))
    xp_assert_close(tf.translation, t, atol=atol)
    # Test that we don't return views that would modify the original array
    xpx.at(tf.translation)[..., 0].set(0.0)
    xp_assert_close(tf.translation, t, atol=atol)
    assert tf.single == (shape == ())


@make_xp_test_case((RigidTransform, "__getitem__"))
def test_indexing(xp):
    atol = 1e-12

    # Test indexing for multiple transforms
    r = Rotation.from_euler('zyx', xp.asarray([[90, 0, 0], [0, 90, 0]]), degrees=True)
    t = xp.asarray([[1.0, 2, 3], [4, 5, 6]])
    tf = RigidTransform.from_components(t, r)

    # Test single index
    xp_assert_close(tf[0].as_matrix()[:3, :3], r[0].as_matrix(), atol=atol)
    xp_assert_close(tf[0].as_matrix()[:3, 3], t[0, ...], atol=atol)

    # Test slice
    tf_slice = tf[0:2]
    xp_assert_close(tf_slice.as_matrix()[:, :3, :3], r[0:2].as_matrix(), atol=atol)
    xp_assert_close(tf_slice.as_matrix()[:, :3, 3], t[0:2, ...], atol=atol)

    # Test boolean indexing
    tf_masked = tf[xp.asarray([True, True])]
    xp_assert_close(tf_masked.as_matrix()[:, :3, :3], r.as_matrix(), atol=atol)
    xp_assert_close(tf_masked.as_matrix()[:, :3, 3], t, atol=atol)

    tf_masked = tf[xp.asarray([False, True])]
    xp_assert_close(tf_masked.as_matrix()[:, :3, :3],
                    r[xp.asarray([False, True])].as_matrix(), atol=atol)
    xp_assert_close(tf_masked.as_matrix()[:, :3, 3], t[xp.asarray([False, True])],
                    atol=atol)

    tf_masked = tf[xp.asarray([False, False])]
    assert len(tf_masked) == 0

    # Test integer array indexing
    idx = xp.asarray([0, 1])
    xp_assert_close(tf[idx].as_matrix()[:, :3, :3], r[idx].as_matrix(), atol=atol)
    xp_assert_close(tf[idx].as_matrix()[:, :3, 3], t, atol=atol)


def test_indexing_array_like():
    atol = 1e-12

    r = Rotation.from_euler('zyx', np.array([[90, 0, 0], [0, 90, 0]]), degrees=True)
    t = np.array([[1.0, 2, 3], [4, 5, 6]])
    tf = RigidTransform.from_components(t, r)

    tf_masked = tf[[False, True]]
    xp_assert_close(tf_masked.as_matrix()[:, :3, :3], r[[False, True]].as_matrix(),
                    atol=atol)
    xp_assert_close(tf_masked.as_matrix()[:, :3, 3], t[[False, True]], atol=atol)
    tf_masked = tf[[False, False]]
    assert len(tf_masked) == 0


@make_xp_test_case((RigidTransform, "concatenate"))
def test_concatenate(xp):
    atol = 1e-12

    # Test concatenation of transforms
    t1 = xp.asarray([1, 0, 0])
    r1 = Rotation.from_euler('z', xp.asarray(90), degrees=True)
    tf1 = RigidTransform.from_components(t1, r1)

    t2 = xp.asarray([0, 1, 0])
    r2 = Rotation.from_euler('x', xp.asarray(90), degrees=True)
    tf2 = RigidTransform.from_components(t2, r2)

    # Concatenate single transforms
    concatenated1 = RigidTransform.concatenate([tf1, tf2])
    xp_assert_close(concatenated1[0].as_matrix(), tf1.as_matrix(), atol=atol)
    xp_assert_close(concatenated1[1].as_matrix(), tf2.as_matrix(), atol=atol)

    # Concatenate multiple transforms
    concatenated2 = RigidTransform.concatenate([tf1, concatenated1])
    xp_assert_close(concatenated2[0].as_matrix(), tf1.as_matrix(), atol=atol)
    xp_assert_close(concatenated2[1].as_matrix(), tf1.as_matrix(), atol=atol)
    xp_assert_close(concatenated2[2].as_matrix(), tf2.as_matrix(), atol=atol)

    # Test ND concatenation
    tf3 = RigidTransform.from_translation(xp.reshape(xp.arange(18), (3, 2, 3)))
    tf4 = RigidTransform.from_translation(xp.reshape(xp.arange(18) + 18, (3, 2, 3)))
    concatenated3 = RigidTransform.concatenate([tf3, tf4])
    xp_assert_close(concatenated3.as_matrix()[:3, ...], tf3.as_matrix(), atol=atol)
    xp_assert_close(concatenated3.as_matrix()[3:, ...], tf4.as_matrix(), atol=atol)


@make_xp_test_case((RigidTransform, "from_matrix"))
def test_input_validation(xp):
    # Test invalid matrix shapes
    inputs = [xp.eye(3), xp.zeros((4, 3)), []]
    for input in inputs:
        with pytest.raises(ValueError, match="Expected `matrix` to have shape"):
            RigidTransform.from_matrix(input)

    # Test invalid last row
    for ndim in range(3):
        shape = (ndim,) * (ndim - 1)
        matrix = xp.zeros(shape + (4, 4))
        matrix = xpx.at(matrix)[...].set(xp.eye(4))
        matrix = xpx.at(matrix)[..., 3, :].set(xp.asarray([1.0, 0, 0, 1]))
        if is_lazy_array(matrix):
            matrix = RigidTransform.from_matrix(matrix).as_matrix()
            assert xp.all(xp.isnan(matrix[..., 3, :]))
        else:
            with pytest.raises(ValueError, match="last row of transformation matrix"):
                RigidTransform.from_matrix(matrix)

    # Test left handed rotation matrix
    matrix = xp.eye(4)
    matrix = xpx.at(matrix)[0, 0].set(-1)
    if is_lazy_array(matrix):
        matrix = RigidTransform.from_matrix(matrix).as_matrix()
        assert xp.all(xp.isnan(matrix[..., :3, :3]))
    else:
        with pytest.raises(ValueError, match="Non-positive determinant"):
            RigidTransform(matrix, normalize=True)

    # Test non-Rotation input
    with pytest.raises(TypeError,
                       match="Expected `rotation` to be a `Rotation` instance"):
        RigidTransform.from_rotation(xp.eye(3))


@make_xp_test_case((RigidTransform, "mean"))
@pytest.mark.parametrize("ndim", range(1, 4))
def test_mean(xp, ndim: int):
    atol = 1e-12
    rng = np.random.default_rng(123)

    dtype = xpx.default_dtype(xp)
    t = xp.asarray(rng.normal(size=(ndim,) * (ndim - 1) + (3,)), dtype=dtype)
    q = xp.asarray(rng.normal(size=(ndim,) * (ndim - 1) + (4,)), dtype=dtype)
    r = Rotation.from_quat(q)
    tf = RigidTransform.from_components(t, r)

    # Unweighted mean
    axis = tuple(range(t.ndim - 1))
    t_mean = xp.mean(t, axis=axis)
    r_mean = r.mean()
    tf_mean = tf.mean()
    assert tf_mean.shape == ()
    xp_assert_close(tf_mean.as_matrix(),
                    RigidTransform.from_components(t_mean, r_mean).as_matrix(),
                    atol=atol)

    # Weighted mean
    if ndim == 1:
        weights = None
        t_mean = t
    else:
        weights = xp.asarray(rng.random(size=(ndim,) * (ndim - 1)), dtype=dtype)
        norm = xp.sum(weights[..., None], axis=axis)
        wsum = xp.sum(t * weights[..., None], axis=axis)
        t_mean = wsum/norm
    r_mean = r.mean(weights=weights)
    tf_mean = tf.mean(weights=weights)
    assert tf_mean.shape == ()
    xp_assert_close(tf_mean.as_matrix(),
                    RigidTransform.from_components(t_mean, r_mean).as_matrix(),
                    atol=atol)


@make_xp_test_case(
    (RigidTransform, "from_rotation"), (RigidTransform, "mean"), (Rotation, "magnitude")
)
@pytest.mark.parametrize("ndim", range(1, 5))
def test_mean_axis(xp, ndim: int):
    axes = xp.tile(xp.concat((-xp.eye(3), xp.eye(3))), (3,) * (ndim - 1) + (1, 1))
    theta = xp.pi / 4
    r = Rotation.from_rotvec(theta * axes)
    tf = RigidTransform.from_rotation(r)

    # Test mean over last axis
    desired = xp.full(axes.shape[:-2], 0.0)
    if ndim == 1:
        desired = desired[()]
    atol = 1e-6 if xpx.default_dtype(xp) is xp.float32 else 1e-10
    xp_assert_close(tf.mean(axis=-1).rotation.magnitude(), desired, atol=atol)

    # Test tuple axes
    desired = xp.full(axes.shape[1:-2], 0.0)
    if ndim < 3:
        desired = desired[()]
    xp_assert_close(tf.mean(axis=(0, -1)).rotation.magnitude(), desired, atol=atol)

    # Empty axis tuple should return RigidTransform unchanged
    tf_mean = tf.mean(axis=())
    xp_assert_close(tf_mean.as_matrix(), tf.as_matrix(), atol=atol)


@make_xp_test_case((RigidTransform, "mean"), (Rotation, "magnitude"))
def test_mean_compare_axis(xp):
    # Create a random set of transforms and compare the mean over an axis with
    # the mean without axis of the sliced transform
    atol = 1e-10 if xpx.default_dtype(xp) == xp.float64 else 1e-6
    rng = np.random.default_rng(0)
    q = xp.asarray(rng.normal(size=(4, 5, 6, 4)), dtype=xpx.default_dtype(xp))
    r = Rotation.from_quat(q)
    t = xp.asarray(rng.normal(size=(4, 5, 6, 3)), dtype=xpx.default_dtype(xp))
    tf = RigidTransform.from_components(t, r)

    mean_0 = tf.mean(axis=0)
    for i in range(q.shape[1]):
        for j in range(q.shape[2]):
            r_slice = Rotation.from_quat(q[:, i, j, ...])
            t_slice = t[:, i, j, ...]
            mean_slice_tf = RigidTransform.from_components(t_slice, r_slice).mean()
            xp_assert_close(
                (mean_0[i][j].rotation * mean_slice_tf.rotation.inv()).magnitude(),
                xp.asarray(0.0)[()], atol=atol,
            )
            xp_assert_close(
                mean_0[i][j].translation, mean_slice_tf.translation, atol=atol,
            )
    mean_1_2 = tf.mean(axis=(1, 2))
    for i in range(q.shape[0]):
        r_slice = Rotation.from_quat(q[i, ...])
        t_slice = t[i, ...]
        mean_slice_tf = RigidTransform.from_components(t_slice, r_slice).mean()
        xp_assert_close(
            (mean_1_2[i].rotation * mean_slice_tf.rotation.inv()).magnitude(),
            xp.asarray(0.0)[()], atol=atol,
        )
        xp_assert_close(
            mean_1_2[i].translation, mean_slice_tf.translation, atol=atol,
        )


@make_xp_test_case((RigidTransform, "mean"))
def test_mean_invalid_weights(xp):
    tf = RigidTransform.from_matrix(xp.tile(xp.eye(4), (4, 1, 1)))
    if is_lazy_array(tf.as_matrix()):
        m = tf.mean(weights=-xp.ones(4))
        assert xp.all(xp.isnan(m.as_matrix()))
    else:
        with pytest.raises(ValueError, match="non-negative"):
            tf.mean(weights=-xp.ones(4))

    # Test weight shape mismatch
    tf = RigidTransform.from_matrix(xp.eye(4))
    with pytest.raises(ValueError, match="Expected `weights` to"):
        tf.mean(weights=xp.ones((2,)))
    tf = RigidTransform.from_matrix(xp.tile(xp.eye(4), (3, 2, 1, 1, 1)))
    with pytest.raises(ValueError, match="Expected `weights` to"):
        tf.mean(weights=xp.ones((2, 1)))


@make_xp_test_case((RigidTransform, "from_translation"))
def test_translation_validation(xp):
    # Test invalid translation shapes
    with pytest.raises(ValueError, match="Expected `translation` to have shape"):
        RigidTransform.from_translation(xp.asarray([1, 2]))

    with pytest.raises(ValueError, match="Expected `translation` to have shape"):
        RigidTransform.from_translation(xp.zeros((2, 2)))


@make_xp_test_case((RigidTransform, "apply"))
def test_vector_validation(xp):
    tf = rigid_transform_to_xp(RigidTransform.identity(2), xp=xp)

    # Test invalid vector shapes
    with pytest.raises(ValueError, match="Expected vector to have shape"):
        tf.apply(xp.asarray([1, 2]))

    with pytest.raises(ValueError, match="Expected vector to have shape"):
        tf.apply(xp.zeros((2, 2)))

    with pytest.raises(ValueError, match="operands could not be broadcast"):
        tf.apply(xp.zeros((1, 4, 3)))


@make_xp_test_case((RigidTransform, "__getitem__"))
def test_indexing_validation(xp):
    tf = RigidTransform.from_matrix(xp.eye(4))

    # Test indexing on single transform
    with pytest.raises(TypeError, match="Single transform is not subscriptable"):
        tf[0]

    with pytest.raises(TypeError, match="Single transform is not subscriptable"):
        tf[0:1]

    # Test length on single transform
    with pytest.raises(TypeError, match="Single transform has no len"):
        len(tf)


@make_xp_test_case((RigidTransform, "__mul__"))
def test_composition_validation(xp):
    tf2 = RigidTransform.from_translation(xp.asarray([[1, 2, 3], [4, 5, 6]]))
    tf3 = RigidTransform.from_translation(xp.asarray([[1, 2, 3], [4, 5, 6], [7, 8, 9]]))

    # Test incompatible shapes
    with pytest.raises(ValueError, match="Cannot broadcast"):
        tf2 * tf3

    tf4 = RigidTransform.from_matrix(xp.tile(xp.eye(4), (1, 4, 1, 1)))
    # Test invalid broadcasting shape
    with pytest.raises(ValueError, match="Cannot broadcast"):
        tf2 * tf4


@make_xp_test_case((RigidTransform, "__mul__"), (RigidTransform, "__rmul__"))
def test_rotation_promotion(xp):
    """Test that Rotation is promoted to RigidTransform in composition."""
    atol = 1e-12
    dtype = xpx.default_dtype(xp)
    rng = np.random.default_rng(0)

    t = xp.asarray(rng.normal(size=(3,)), dtype=dtype)
    r1 = rotation_to_xp(Rotation.random(rng=rng), xp)
    r2 = rotation_to_xp(Rotation.random(rng=rng), xp)
    tf = RigidTransform.from_components(t, r2)

    # RigidTransform * Rotation
    result = tf * r1
    expected = tf * RigidTransform.from_rotation(r1)
    assert isinstance(result, RigidTransform)
    xp_assert_close(result.as_matrix(), expected.as_matrix(), atol=atol)

    # Rotation * RigidTransform
    result = r1 * tf
    expected = RigidTransform.from_rotation(r1) * tf
    assert isinstance(result, RigidTransform)
    xp_assert_close(result.as_matrix(), expected.as_matrix(), atol=atol)


@make_xp_test_case((RigidTransform, "concatenate"))
def test_concatenate_validation(xp):
    tf = RigidTransform.from_matrix(xp.eye(4))

    # Test invalid inputs
    with pytest.raises(TypeError,
                       match="input must contain RigidTransform objects"):
        RigidTransform.concatenate([tf, xp.eye(4)])

    # Test incompatible shapes
    tf2 = RigidTransform.from_translation(xp.ones((1, 1, 1, 3)))
    # Frameworks have a highly heterogeneous way of reporting errors for this case
    with pytest.raises((ValueError, TypeError, RuntimeError)):
        RigidTransform.concatenate([tf, tf2])


@make_xp_test_case((RigidTransform, "__setitem__"))
def test_setitem(xp):
    tf = RigidTransform.from_translation(xp.asarray([[1, 2, 3], [4, 5, 6], [7, 8, 9]]))
    single = RigidTransform.from_translation(xp.asarray([1, 1, 1]))
    double = RigidTransform.from_translation(xp.asarray([[2, 2, 2], [3, 3, 3]]))
    triple = RigidTransform.from_translation(xp.asarray([[3, 3, 3],
                                                         [4, 4, 4],
                                                         [5, 5, 5]]))

    # Test indexing with integer index
    tf[0] = single
    xp_assert_close(tf.translation, xp.asarray([[1.0, 1, 1], [4, 5, 6], [7, 8, 9]]))

    # Test indexing with slice
    tf = RigidTransform.from_translation(xp.asarray([[1, 2, 3], [4, 5, 6], [7, 8, 9]]))
    tf[:2] = double
    xp_assert_close(tf.translation, xp.asarray([[2.0, 2, 2], [3, 3, 3], [7, 8, 9]]))

    # Test indexing with ellipsis
    tf = RigidTransform.from_translation(xp.asarray([[1, 2, 3], [4, 5, 6], [7, 8, 9]]))
    tf[...] = triple
    xp_assert_close(tf.translation, xp.asarray([[3.0, 3, 3], [4, 4, 4], [5, 5, 5]]))

    # Test indexing with boolean array
    tf = RigidTransform.from_translation(xp.asarray([[1, 2, 3], [4, 5, 6], [7, 8, 9]]))
    mask = xp.asarray([True, False, True])
    tf[mask] = double
    xp_assert_close(tf.translation, xp.asarray([[2.0, 2, 2], [4, 5, 6], [3, 3, 3]]))


@make_xp_test_case((RigidTransform, "__setitem__"))
@pytest.mark.skip_xp_backends("array_api_strict",
                              reason="doesn't support fancy indexing __setitem__")
def test_setitem_fancy_indexing(xp):
    double = RigidTransform.from_translation(xp.asarray([[2, 2, 2], [3, 3, 3]]))
    tf = RigidTransform.from_translation(xp.asarray([[1, 2, 3], [4, 5, 6], [7, 8, 9]]))
    idx = xp.asarray([0, 2])
    tf[idx] = double
    xp_assert_close(tf.translation, xp.asarray([[2.0, 2, 2], [4, 5, 6], [3, 3, 3]]))


@make_xp_test_case((RigidTransform, "__setitem__"))
def test_setitem_validation(xp):
    tf = RigidTransform.from_translation(xp.asarray([[1, 2, 3], [4, 5, 6]]))
    single = RigidTransform.from_matrix(xp.eye(4))

    # Test setting item on single transform
    with pytest.raises(TypeError, match="Single transform is not subscriptable"):
        single[0] = tf

    # Test invalid value type
    with pytest.raises(TypeError, match="value must be a RigidTransform"):
        tf[0] = xp.eye(4)


@pytest.mark.skip_xp_backends("jax.numpy",
                              reason="JAX does not support memory sharing")
@make_xp_test_case((RigidTransform, "as_matrix"))
def test_copy_flag(xp):
    # Test that copy=True creates new memory
    matrix = xp.eye(4)
    tf = RigidTransform(matrix, normalize=False, copy=True)
    matrix[0, 0] = 2
    assert tf.as_matrix()[0, 0] == 1

    # Test that copy=False shares memory
    matrix = xp.eye(4)
    tf = RigidTransform(matrix, normalize=False, copy=False)
    matrix[0, 0] = 2
    assert tf.as_matrix()[0, 0] == 2


@make_xp_test_case(normalize_dual_quaternion)
@pytest.mark.parametrize("ndim", range(1, 4))
def test_normalize_dual_quaternion(xp, ndim: int):
    dtype = xpx.default_dtype(xp)
    atol = 1e-12 if dtype == xp.float64 else 1e-6
    rng = np.random.default_rng(100)
    shape = (ndim,) * (ndim - 1)

    dual_quat = normalize_dual_quaternion(xp.zeros((1, 8)))
    xp_assert_close(xp_vector_norm(dual_quat[0, :4], axis=-1), xp.asarray(1.0)[()],
                    atol=1e-12)
    xp_assert_close(xp.vecdot(dual_quat[0, :4], dual_quat[0, 4:])[()],
                    xp.asarray(0.0)[()], atol=1e-12)

    dual_quat = xp.asarray(rng.normal(size=shape + (8,)), dtype=dtype)
    dual_quat = normalize_dual_quaternion(dual_quat)
    expected = xp.ones(shape) if shape != () else xp.asarray(1.0)[()]
    xp_assert_close(xp_vector_norm(dual_quat[..., :4], axis=-1), expected, atol=atol)
    expected = xp.zeros(shape) if shape != () else xp.asarray(0.0)[()]
    vecdot = xp.vecdot(dual_quat[..., :4], dual_quat[..., 4:])
    vecdot = vecdot[()] if vecdot.shape == () else vecdot
    xp_assert_close(vecdot, expected, atol=atol)


@make_xp_test_case(
    (RigidTransform, "from_matrix"), (RigidTransform, "from_rotation"),
    (RigidTransform, "from_translation"), (RigidTransform, "from_components"),
    (RigidTransform, "from_exp_coords"), (RigidTransform, "from_dual_quat")
)
def test_empty_transform_construction(xp):
    tf = RigidTransform.from_matrix(xp.empty((0, 4, 4)))
    assert len(tf) == 0
    assert not tf.single

    tf = RigidTransform.from_rotation(Rotation.from_quat(xp.zeros((0, 4))))
    assert len(tf) == 0
    assert not tf.single

    tf = RigidTransform.from_translation(xp.empty((0, 3)))
    assert len(tf) == 0
    assert not tf.single

    empty_rot = Rotation.from_quat(xp.zeros((0, 4)))
    tf = RigidTransform.from_components(xp.empty((0, 3)), empty_rot)
    assert len(tf) == 0
    assert not tf.single

    tf = RigidTransform.from_exp_coords(xp.empty((0, 6)))
    assert len(tf) == 0
    assert not tf.single

    tf = RigidTransform.from_dual_quat(xp.empty((0, 8)))
    assert len(tf) == 0
    assert not tf.single

    tf = RigidTransform.identity(0)
    assert len(tf) == 0
    assert not tf.single


@make_xp_test_case((RigidTransform, "from_matrix"), (RigidTransform, "as_components"),
                   (RigidTransform, "as_exp_coords"), (RigidTransform, "as_dual_quat"))
def test_empty_transform_representation(xp):
    tf = RigidTransform.from_matrix(xp.empty((0, 4, 4)))

    assert len(tf.rotation) == 0
    assert tf.translation.shape == (0, 3)

    t, r = tf.as_components()
    assert t.shape == (0, 3)
    assert len(r) == 0

    assert tf.as_matrix().shape == (0, 4, 4)
    assert tf.as_exp_coords().shape == (0, 6)
    assert tf.as_dual_quat().shape == (0, 8)


@make_xp_test_case((RigidTransform, "from_matrix"), (RigidTransform, "apply"))
def test_empty_transform_application(xp):
    tf = RigidTransform.from_matrix(xp.empty((0, 4, 4)))

    assert tf.apply(xp.zeros((3,))).shape == (0, 3)
    assert tf.apply(xp.empty((0, 3))).shape == (0, 3)

    with pytest.raises(ValueError, match="operands could not be broadcast together"):
        tf.apply(xp.zeros((2, 3)))


@make_xp_test_case((RigidTransform, "from_matrix"), (RigidTransform, "__mul__"))
def test_empty_transform_composition(xp):
    tf_empty = RigidTransform.from_matrix(xp.empty((0, 4, 4)))
    tf_single = RigidTransform.from_matrix(xp.eye(4))
    tf_many = rigid_transform_to_xp(RigidTransform.identity(3), xp=xp)

    assert len(tf_empty * tf_empty) == 0
    assert len(tf_empty * tf_single) == 0
    assert len(tf_single * tf_empty) == 0

    with pytest.raises(ValueError, match="Cannot broadcast"):
        tf_many * tf_empty

    with pytest.raises(ValueError, match="Cannot broadcast"):
        tf_empty * tf_many


@make_xp_test_case((RigidTransform, "from_matrix"), (RigidTransform, "concatenate"))
def test_empty_transform_concatenation(xp):
    tf_empty = RigidTransform.from_matrix(xp.empty((0, 4, 4)))
    tf_single = RigidTransform.from_matrix(xp.eye(4))
    tf_many = rigid_transform_to_xp(RigidTransform.identity(2), xp=xp)

    assert len(RigidTransform.concatenate([tf_empty, tf_empty])) == 0
    assert len(RigidTransform.concatenate([tf_empty, tf_single])) == 1
    assert len(RigidTransform.concatenate([tf_single, tf_empty])) == 1
    assert len(RigidTransform.concatenate([tf_empty, tf_many])) == 2
    assert len(RigidTransform.concatenate([tf_many, tf_empty])) == 2
    assert len(RigidTransform.concatenate([tf_many, tf_empty, tf_single])) == 3


@make_xp_test_case((RigidTransform, "from_matrix"), (RigidTransform, "inv"),
                   (RigidTransform, "__pow__"))
def test_empty_transform_inv_and_pow(xp):
    tf = RigidTransform.from_matrix(xp.empty((0, 4, 4)))
    assert len(tf.inv()) == 0
    assert len(tf ** 0) == 0
    assert len(tf ** 1) == 0
    assert len(tf ** -1) == 0
    assert len(tf ** 0.5) == 0


@make_xp_test_case((RigidTransform, "__getitem__"))
def test_empty_transform_indexing(xp):
    tf_many = rigid_transform_to_xp(RigidTransform.identity(3), xp=xp)
    tf_zero = tf_many[xp.asarray([], dtype=xp.int32)]
    assert len(tf_zero) == 0

    # Array API does not specify out-of-bounds indexing. Only check for numpy.
    if is_numpy(xp):
        assert len(tf_zero[:5]) == 0  # Slices can go out of bounds.

    with pytest.raises(IndexError):
        tf_zero[0]

    with pytest.raises(IndexError):
        tf_zero[xp.asarray([0, 2])]

    with pytest.raises(IndexError):
        tf_zero[xp.asarray([False, True])]


@make_xp_test_case((RigidTransform, "from_matrix"))
@pytest.mark.skip_xp_backends("array_api_strict",
                              reason="array API doesn't support pickling")
def test_pickling(xp):
    # Note: Array API makes no provision for arrays to be pickleable, so
    # it's OK to skip this test for the backends that don't support it
    mat = xp.eye(4)
    mat = xpx.at(mat)[0, 3].set(2.0)
    tf = RigidTransform.from_matrix(mat)
    pkl = pickle.dumps(tf)
    unpickled = pickle.loads(pkl)
    xp_assert_close(tf.as_matrix(), unpickled.as_matrix(), atol=1e-15)


@make_xp_test_case(
    (RigidTransform, "as_matrix"),
    (RigidTransform, "__iter__"),
    (RigidTransform, "identity"),
)
def test_rigid_transform_iter(xp):
    r = rigid_transform_to_xp(RigidTransform.identity(3), xp)
    for i, r_i in enumerate(r):
        assert isinstance(r_i, RigidTransform)
        xp_assert_equal(r_i.as_matrix(), r[i].as_matrix())
        if i > len(r):
            raise RuntimeError("Iteration exceeded length of transforms")


@make_xp_test_case((RigidTransform, "from_translation"))
@pytest.mark.parametrize("dim", range(1, 5))
def test_shape_property(xp, dim: int):
    shape = (dim,) * (dim - 1)
    tf = RigidTransform.from_translation(xp.zeros(shape + (3,)))
    assert tf.shape == shape


def test_non_writeable():
    mat = np.eye(4)
    mat.flags.writeable = False
    RigidTransform.from_matrix(mat)  # Regression test against gh-24378
