351class BaseTransform(ABC):
352    """
353    Abstract interface for all spatial transformations.
354    """
355
356    def __init__(self, dtype: np.dtype = np.float64):
357        self.dtype = dtype
358
359    @abstractmethod
360    def as_matrix() -> np.ndarray:
361        """
362        Returns the 4x4 homogeneous representation of the transform.
363
364        Returns:
365            np.ndarray: 4x4 matrix of the transform's dtype.
366        """
367        pass
368
369    @abstractmethod
370    def inverse() -> BaseTransform:
371        """
372        Returns the mathematical inverse of the transformation.
373
374        Returns:
375            BaseTransform: The inverse transformation.
376        """
377        pass
378
379    def _apply(self, vector: np.ndarray | list | tuple) -> np.ndarray:
380        """
381        Apply the transform to 3D vectors (Nx3 or Nx4).
382
383        Args:
384            vector: Nx3 or Nx4 array of vectors.
385                    - If Nx3: Treated as 3D points/vectors (w=1 implied if Transform,
386                      checking subclass logic).
387                      Standard BaseTransform behavior:
388                      Homogenize (w=1) -> Multiply -> Dehomogenize (w division).
389                    - If Nx4: Treated as homogeneous vectors.
390                      Multiply -> Return Nx4.
391
392        Returns:
393            np.ndarray: Transformed vectors (Nx3 or Nx4).
394        """
395        vector = np.atleast_2d(vector)
396
397        if vector.shape[1] == 3:
398            # Homogenize (w=1)
399            hom_vector = np.hstack([vector, np.ones((vector.shape[0], 1), dtype=self.dtype)])
400            # Apply
401            transformed_hom = (self.as_matrix() @ hom_vector.T).T
402            # Dehomogenize (return 3D)
403            return transformed_hom[:, :3]
404
405        elif vector.shape[1] == 4:
406            # Generic 4x4
407            return (self.as_matrix() @ vector.T).T
408        else:
409            raise ValueError(f"Input vector must be Nx3 or Nx4, got {vector.shape}")
410
411    @abstractmethod
412    def __mul__(self, other: BaseTransform) -> BaseTransform:
413        """
414        Composes this transform with another.
415        Composition follows standard matrix multiplication order: (T1 * T2) * p = T1 * (T2 * p).
416
417        Args:
418            other: The transform to apply second.
419
420        Returns:
421            BaseTransform: The composed transformation.
422        """
423        pass
424
425    @abstractmethod
426    def to_dict(self) -> dict[str, Any]:
427        """
428        Serialize the transform to a JSON-compatible dictionary.
429
430        The dictionary MUST include a "type" key with the class name
431        to enable proper deserialization.
432
433        Returns:
434            Dict[str, Any]: JSON-compatible dictionary representation.
435        """
436        pass
437
438    @classmethod
439    @abstractmethod
440    def from_dict(cls, data: dict[str, Any]) -> BaseTransform:
441        """
442        Deserialize a transform from a dictionary.
443
444        Args:
445            data: Dictionary previously created by to_dict().
446
447        Returns:
448            BaseTransform: The deserialized transform instance.
449        """
450        pass
451
452    def __repr__(self) -> str:
453        return f"<{self.__class__.__name__}>"

456@register_transform
457class Transform(BaseTransform):
458    """
459    Standard SE(3) rigid body transformation.
460    Consists of a translation (3x1) and a rotation (quaternion).
461    """
462
463    def __init__(
464        self,
465        translation: np.ndarray | list | tuple | None = None,
466        rotation: quaternion.quaternion | np.ndarray | list | tuple | Transform | None = None,
467        dtype: np.dtype = np.float64,
468    ):
469        super().__init__(dtype=dtype)
470
471        self.translation = ensure_translation(translation, self.dtype)
472
473        # Handle rotation: can be quaternion, array, or Transform/Rotation object
474        if isinstance(rotation, BaseTransform):
475            # It's a Transform or Rotation object - extract the quaternion
476            rotation = rotation.rotation
477
478        self.rotation = ensure_rotation(rotation, self.dtype)
479
480    @classmethod
481    def from_matrix(cls, matrix: np.ndarray, dtype: np.dtype | None = None) -> Transform:
482        """
483        Creates a Transform from a 4x4 matrix.
484        """
485        if matrix.shape != (4, 4):
486            raise ValueError(f"Matrix must be 4x4, got {matrix.shape}")
487
488        target_dtype = dtype if dtype is not None else matrix.dtype
489        translation = matrix[:3, 3]
490        rot_mat = matrix[:3, :3]
491        rot_quat = quaternion.from_rotation_matrix(rot_mat)
492        return cls(translation=translation, rotation=rot_quat, dtype=target_dtype)
493
494    @classmethod
495    def from_rotation_matrix(
496        cls,
497        rotation_matrix: np.ndarray,
498        t: np.ndarray | list | tuple | None = None,
499        *,
500        validate: bool = True,
501        dtype: np.dtype = np.float64,
502    ) -> Transform:
503        """Create a Transform from a 3x3 rotation matrix.
504
505        Args:
506            rotation_matrix: A 3x3 rotation matrix (must be SO(3) when
507                ``validate=True``).
508            t: Optional 3-element translation vector.
509            validate: If True (default), verify that ``rotation_matrix`` is a
510                valid SO(3) element (orthogonal with determinant +1). Set to
511                False to skip the check on trusted data. Follows the same
512                pattern as ``scipy.spatial.transform.Rotation.from_matrix``
513                with ``assume_valid``.
514            dtype: NumPy dtype for the resulting transform.
515
516        Returns:
517            Transform: The constructed SE(3) transform.
518
519        Raises:
520            ValueError: If the matrix is not 3x3 or fails SO(3) validation.
521
522        Example:
523            >>> R = np.eye(3)
524            >>> tf.Transform.from_rotation_matrix(R, t=[1, 2, 3])
525        """
526        R = np.asarray(rotation_matrix, dtype=np.float64)
527        if R.shape != (3, 3):
528            raise ValueError(f"Rotation matrix must be 3x3, got {R.shape}")
529
530        if validate:
531            # Check orthogonality: R^T R ≈ I
532            orthogonality_error = np.max(np.abs(R.T @ R - np.eye(3)))
533            if orthogonality_error > 1e-6:
534                raise ValueError(
535                    f"Rotation matrix is not orthogonal (max error {orthogonality_error:.2e}). "
536                    "Expected R^T R = I. Pass validate=False to skip this check."
537                )
538            # Check determinant ≈ +1 (proper rotation, not reflection)
539            det = np.linalg.det(R)
540            if abs(det - 1.0) > 1e-6:
541                raise ValueError(
542                    f"Rotation matrix determinant is {det:.6f}, expected +1.0 (SO(3)). "
543                    "A determinant of -1 indicates a reflection, not a rotation."
544                )
545
546        rot_quat = quaternion.from_rotation_matrix(R)
547        return cls(translation=t, rotation=rot_quat, dtype=dtype)
548
549    @classmethod
550    def from_quaternion(
551        cls,
552        q: quaternion.quaternion | np.ndarray | list | tuple,
553        t: np.ndarray | list | tuple | None = None,
554        *,
555        convention: str = "wxyz",
556        dtype: np.dtype = np.float64,
557    ) -> Transform:
558        """Create a Transform from a quaternion.
559
560        The quaternion is auto-normalized to unit length.
561
562        Args:
563            q: Quaternion — either a ``quaternion.quaternion`` object or a
564                4-element array. The element order is determined by
565                ``convention``.
566            t: Optional 3-element translation vector.
567            convention: Element ordering of array input:
568                ``"wxyz"`` (default, numpy-quaternion / Hamilton) or
569                ``"xyzw"`` (scipy / ROS). Ignored when ``q`` is a
570                ``quaternion.quaternion`` object.
571            dtype: NumPy dtype for the resulting transform.
572
573        Returns:
574            Transform: The constructed SE(3) transform.
575
576        Raises:
577            ValueError: If the quaternion is zero or the convention is unknown.
578
579        Example:
580            >>> tf.Transform.from_quaternion([1, 0, 0, 0])  # wxyz identity
581            >>> tf.Transform.from_quaternion([0, 0, 0, 1], convention="xyzw")
582        """
583        if isinstance(q, quaternion.quaternion):
584            quat = q
585        else:
586            arr = np.asarray(q, dtype=np.float64).ravel()
587            if arr.size != 4:
588                raise ValueError(f"Quaternion must have 4 elements, got {arr.size}")
589            if convention == "wxyz":
590                quat = quaternion.quaternion(*arr)
591            elif convention == "xyzw":
592                x, y, z, w = arr
593                quat = quaternion.quaternion(w, x, y, z)
594            else:
595                raise ValueError(
596                    f"Unknown quaternion convention '{convention}'. "
597                    "Use 'wxyz' (numpy-quaternion) or 'xyzw' (scipy/ROS)."
598                )
599
600        # Normalize — reject zero
601        norm = np.sqrt(quat.w**2 + quat.x**2 + quat.y**2 + quat.z**2)
602        if norm < 1e-15:
603            raise ValueError("Zero-norm quaternion cannot represent a rotation.")
604        quat = quat / norm
605
606        return cls(translation=t, rotation=quat, dtype=dtype)
607
608    @classmethod
609    def from_axis_angle(
610        cls,
611        axis: np.ndarray | list | tuple,
612        angle: float,
613        t: np.ndarray | list | tuple | None = None,
614        *,
615        dtype: np.dtype = np.float64,
616    ) -> Transform:
617        """Create a Transform from an axis-angle representation.
618
619        The axis is auto-normalized to unit length.
620
621        Args:
622            axis: 3-element rotation axis (will be normalized).
623            angle: Rotation angle in radians.
624            t: Optional 3-element translation vector.
625            dtype: NumPy dtype for the resulting transform.
626
627        Returns:
628            Transform: The constructed SE(3) transform.
629
630        Raises:
631            ValueError: If the axis has zero length.
632
633        Example:
634            >>> tf.Transform.from_axis_angle([0, 0, 1], np.pi / 2)
635        """
636        axis_array = np.asarray(axis, dtype=np.float64).ravel()
637        if axis_array.size != 3:
638            raise ValueError(f"Axis must have 3 elements, got {axis_array.size}")
639
640        axis_norm = np.linalg.norm(axis_array)
641        if axis_norm < 1e-15:
642            raise ValueError("Zero-length axis vector cannot define a rotation direction.")
643
644        axis_unit = axis_array / axis_norm
645        rotation_vector = axis_unit * float(angle)
646        rot_quat = quaternion.from_rotation_vector(rotation_vector)
647
648        return cls(translation=t, rotation=rot_quat, dtype=dtype)
649
650    def as_matrix(self) -> np.ndarray:
651        """Return the 4x4 homogeneous transformation matrix [R|t; 0 1]."""
652        matrix = np.eye(4, dtype=self.dtype)
653        matrix[:3, :3] = quaternion.as_rotation_matrix(self.rotation).astype(self.dtype)
654        matrix[:3, 3] = self.translation.ravel()
655        return matrix
656
657    def inverse(self) -> Transform:
658        """Return the inverse SE(3) transform: T^-1 = [-R^T t; R^T]."""
659        inv_rotation = self.rotation.conjugate()
660        # Rotate expects (..., 3) vectors. flatten() ensures we get (3,) if translation is (3, 1)
661        inv_translation = -quaternion.rotate_vectors(inv_rotation, self.translation.flatten())
662        return Transform(translation=inv_translation, rotation=inv_rotation, dtype=self.dtype)
663
664    def __mul__(self, other: BaseTransform) -> BaseTransform:
665        if isinstance(other, Transform):
666            # T1 * T2 = [R1*R2, R1*t2 + t1]
667            new_rotation = self.rotation * other.rotation
668            new_translation = (
669                quaternion.rotate_vectors(self.rotation, other.translation.flatten())
670                + self.translation.flatten()
671            )
672            return Transform(translation=new_translation, rotation=new_rotation, dtype=self.dtype)
673
674        if isinstance(other, (Projection, CameraProjection)):
675            # Transform * Projection -> Invalid
676            # 3D->3D * 3D->2D = dimensional mismatch if interpreted strictly as flow?
677            # Guidelines say: Transform * CameraProjection = Forbidden.
678            raise TypeError(
679                "Composition 'Transform * CameraProjection' is invalid. "
680                "Transforms (Spatial) cannot pre-multiply Projections (Sensor). "
681                "Did you mean 'CameraProjection * Transform'?"
682            )
683
684        if isinstance(other, InverseProjection):
685            # Transform * InverseProjection -> InverseCompositeProjection
686            # T * P_inv
687            # Logic: P_inv unprojects to frame A. T transforms A to B. Result unprojects to B.
688
689            # If other is already InverseCompositeProjection: T * (T_old * K_inv)
690            # The __rmul__ of InverseCompositeProjection should handle
691            # this if we return NotImplemented?
692            # But BaseTransform doesn't implement __rmul__ universally.
693            # Let's handle it explicitly or rely on commutation if possible.
694            # Transform doesn't know about InverseCompositeProjection class structure necessarily?
695            # But we just added it to this module.
696
697            if isinstance(other, InverseCompositeProjection):
698                # T * ICA(T_old, K_inv) = ICA(T * T_old, K_inv)
699                new_transform = self * other.transform
700                if isinstance(new_transform, Transform):
701                    return InverseCompositeProjection(
702                        new_transform, other.projection, dtype=self.dtype
703                    )
704
705            # Generic InverseProjection (e.g. InverseCameraProjection or raw InverseProjection)
706            # Treat 'self' as the extrinsics T.
707            # Result = T * P_inv.
708            # We create InverseCompositeProjection(T, P_inv).
709            return InverseCompositeProjection(self, other, dtype=self.dtype)
710
711        # Fallback to matrix multiplication
712        return MatrixTransform(self.as_matrix() @ other.as_matrix(), dtype=self.dtype)
713
714    def __repr__(self) -> str:
715        return f"Transform(translation={self.translation.flatten()!r}, rotation={self.rotation!r})"
716
717    def to_dict(self) -> dict[str, Any]:
718        """Serialize transform to a JSON-compatible dictionary."""
719        t = self.translation.flatten()
720        q = self.rotation
721        return {
722            "type": "Transform",
723            "translation": [float(t[0]), float(t[1]), float(t[2])],
724            "rotation": [float(q.w), float(q.x), float(q.y), float(q.z)],
725            "dtype": np.dtype(self.dtype).name,
726        }
727
728    @classmethod
729    def from_dict(cls, data: dict[str, Any]) -> Transform:
730        """Deserialize transform from a dictionary."""
731        dtype = np.dtype(data.get("dtype", "float64"))
732        return cls(
733            translation=data["translation"],
734            rotation=data["rotation"],
735            dtype=dtype,
736        )

739class Translation(Transform):
740    """A Transform with only translation (Identity rotation)."""
741
742    def __init__(
743        self,
744        x: float = 0.0,
745        y: float = 0.0,
746        z: float = 0.0,
747        translation: np.ndarray | list | tuple | None = None,
748    ):
749        if translation is not None:
750            super().__init__(translation=translation)
751        else:
752            super().__init__(translation=[x, y, z])

755class Rotation(Transform):
756    """
757    A Transform with only rotation (Zero translation).
758
759    Supports multiple construction patterns:
760    - Quaternion components: Rotation(w=1, x=0, y=0, z=0)
761    - Quaternion object: Rotation(rotation=q)
762    - Euler angles: Rotation.from_roll_pitch_yaw(roll=0, pitch=0, yaw=0)
763    ...
764    """
765
766    def __init__(
767        self,
768        w: float = 1.0,
769        x: float = 0.0,
770        y: float = 0.0,
771        z: float = 0.0,
772        rotation: quaternion.quaternion | np.ndarray | list | tuple | None = None,
773    ):
774        if rotation is not None:
775            super().__init__(rotation=rotation)
776        else:
777            super().__init__(rotation=[w, x, y, z])
778
779    @classmethod
780    def from_roll_pitch_yaw(
781        cls,
782        roll: float = 0.0,
783        pitch: float = 0.0,
784        yaw: float = 0.0,
785    ) -> Rotation:
786        """
787        Create a Rotation from roll-pitch-yaw angles.
788
789        Uses the aerospace/robotics intrinsic **ZYX** (Tait-Bryan) convention:
790        yaw (Z) → pitch (Y) → roll (X).
791
792        Args:
793            roll: Rotation about X-axis in radians.
794            pitch: Rotation about Y-axis in radians.
795            yaw: Rotation about Z-axis in radians.
796
797        Returns:
798            Rotation: A rotation-only transform.
799
800        Example:
801            >>> attitude = tf.Rotation.from_roll_pitch_yaw(pitch=np.radians(10))
802            >>> heading = tf.Rotation.from_roll_pitch_yaw(yaw=np.pi/4)
803        """
804        q = from_roll_pitch_yaw(roll=roll, pitch=pitch, yaw=yaw)
805        return cls(rotation=q)
806
807    def as_roll_pitch_yaw(self) -> tuple[float, float, float]:
808        """
809        Extract roll, pitch, yaw from the rotation.
810
811        Uses the aerospace/robotics intrinsic **ZYX** (Tait-Bryan) convention.
812
813        Returns:
814            Tuple[float, float, float]: ``(roll, pitch, yaw)`` in radians.
815
816        Warning:
817            Euler angles have a singularity (gimbal lock) when pitch = ±90°.
818
819        Example:
820            >>> rotation = tf.Rotation.from_roll_pitch_yaw(roll=0.1, pitch=0.2, yaw=0.3)
821            >>> roll, pitch, yaw = rotation.as_roll_pitch_yaw()
822        """
823        return as_roll_pitch_yaw(self.rotation)
824
825    @classmethod
826    def from_rotation_matrix(
827        cls,
828        rotation_matrix: np.ndarray,
829        *,
830        validate: bool = True,
831    ) -> Rotation:
832        """Create a Rotation from a 3x3 rotation matrix.
833
834        See :meth:`Transform.from_rotation_matrix` for full documentation.
835
836        Args:
837            rotation_matrix: A 3x3 rotation matrix.
838            validate: If True, verify SO(3) membership.
839
840        Returns:
841            Rotation: A rotation-only transform.
842        """
843        transform = Transform.from_rotation_matrix(rotation_matrix, validate=validate)
844        return cls(rotation=transform.rotation)
845
846    @classmethod
847    def from_quaternion(
848        cls,
849        q: quaternion.quaternion | np.ndarray | list | tuple,
850        *,
851        convention: str = "wxyz",
852    ) -> Rotation:
853        """Create a Rotation from a quaternion.
854
855        See :meth:`Transform.from_quaternion` for full documentation.
856
857        Args:
858            q: Quaternion (object or 4-element array).
859            convention: ``"wxyz"`` or ``"xyzw"``.
860
861        Returns:
862            Rotation: A rotation-only transform.
863        """
864        transform = Transform.from_quaternion(q, convention=convention)
865        return cls(rotation=transform.rotation)
866
867    @classmethod
868    def from_axis_angle(
869        cls,
870        axis: np.ndarray | list | tuple,
871        angle: float,
872    ) -> Rotation:
873        """Create a Rotation from an axis-angle representation.
874
875        See :meth:`Transform.from_axis_angle` for full documentation.
876
877        Args:
878            axis: 3-element rotation axis (auto-normalized).
879            angle: Rotation angle in radians.
880
881        Returns:
882            Rotation: A rotation-only transform.
883        """
884        transform = Transform.from_axis_angle(axis, angle)
885        return cls(rotation=transform.rotation)