Source code for simvx.core.math.transforms

"""Transform2D, Basis, and Transform3D — hierarchical placement types."""

from __future__ import annotations

import math

import numpy as np

from .types import Quat, Vec3

# ============================================================================
# Transform2D
# ============================================================================



class Transform2D:
    """2D transform: position + rotation + scale with cached 3x3 matrix."""

    def __init__(self, position=(0, 0), rotation=0.0, scale=(1, 1)):
        self._position = np.array(position, dtype=np.float32)
        self._rotation = float(rotation)
        self._scale = np.array(scale, dtype=np.float32)
        self._matrix: np.ndarray | None = None
        self._dirty = True

    @property
    def position(self) -> np.ndarray:
        return self._position

    @position.setter
    def position(self, value):
        self._position = np.array(value, dtype=np.float32)
        self._dirty = True

    @property
    def rotation(self) -> float:
        return self._rotation

    @rotation.setter
    def rotation(self, value: float):
        self._rotation = float(value)
        self._dirty = True

    @property
    def scale(self) -> np.ndarray:
        return self._scale

    @scale.setter
    def scale(self, value):
        self._scale = np.array(value, dtype=np.float32)
        self._dirty = True

    def _rebuild_matrix(self):
        c, s = math.cos(self._rotation), math.sin(self._rotation)
        sx, sy = self._scale
        self._matrix = np.array(
            [
                [c * sx, -s * sy, self._position[0]],
                [s * sx, c * sy, self._position[1]],
                [0, 0, 1],
            ],
            dtype=np.float32,
        )

    @property
    def matrix(self) -> np.ndarray:
        if self._dirty:
            self._rebuild_matrix()
            self._dirty = False
        return self._matrix



    def translated(self, offset) -> Transform2D:
        """Return a new transform with position shifted by offset."""
        return Transform2D(self._position + np.array(offset, dtype=np.float32), self._rotation, self._scale)




    def rotated(self, angle: float) -> Transform2D:
        """Return a new transform with additional rotation (radians)."""
        return Transform2D(self._position, self._rotation + angle, self._scale)




    def scaled(self, factor) -> Transform2D:
        """Return a new transform with scale multiplied by factor."""
        if isinstance(factor, int | float):
            f = np.array([factor, factor], dtype=np.float32)
        else:
            f = np.array(factor, dtype=np.float32)
        return Transform2D(self._position, self._rotation, self._scale * f)




    def transform_point(self, point) -> np.ndarray:
        """Transform a 2D point through this transform."""
        p = np.array(point, dtype=np.float32)
        h = np.array([p[0], p[1], 1.0], dtype=np.float32)
        result = self.matrix @ h
        return result[:2]




    def inverse(self) -> Transform2D:
        """Return the inverse transform."""
        inv_s = np.where(np.abs(self._scale) > 1e-10, 1.0 / self._scale, np.zeros_like(self._scale))
        inv_r = -self._rotation
        c, s = math.cos(inv_r), math.sin(inv_r)
        rot = np.array([[c, -s], [s, c]], dtype=np.float32)
        inv_p = -(rot @ (inv_s * self._position))
        return Transform2D(inv_p, inv_r, inv_s)




    def __mul__(self, other):
        if isinstance(other, Transform2D):
            # Compose transforms via matrix multiplication
            m = self.matrix @ other.matrix
            # Extract components from composed matrix
            sx = math.sqrt(m[0, 0] ** 2 + m[1, 0] ** 2)
            sy = math.sqrt(m[0, 1] ** 2 + m[1, 1] ** 2)
            rot = math.atan2(m[1, 0], m[0, 0])
            return Transform2D((m[0, 2], m[1, 2]), rot, (sx, sy))
        if isinstance(other, tuple | list | np.ndarray):
            return self.transform_point(other)
        return NotImplemented




    def __eq__(self, other):
        if isinstance(other, Transform2D):
            return (
                np.allclose(self._position, other._position)
                and abs(self._rotation - other._rotation) < 1e-6
                and np.allclose(self._scale, other._scale)
            )
        return NotImplemented




    def __repr__(self):
        return f"Transform2D(pos={self._position}, rot={self._rotation:.4g}, scale={self._scale})"



# ============================================================================
# Basis
# ============================================================================



class Basis:
    """3x3 rotation/scale matrix. Column access for right/up/forward."""

    __slots__ = ("_m",)

    def __init__(self, matrix: np.ndarray | None = None):
        self._m = matrix.astype(np.float32) if matrix is not None else np.eye(3, dtype=np.float32)

    @property
    def x(self) -> np.ndarray:
        """Right vector (first column)."""
        return self._m[:, 0].copy()

    @property
    def y(self) -> np.ndarray:
        """Up vector (second column)."""
        return self._m[:, 1].copy()

    @property
    def z(self) -> np.ndarray:
        """Forward vector (third column)."""
        return self._m[:, 2].copy()



    def rotated(self, axis, angle: float) -> Basis:
        """Return a new basis rotated around axis by angle (radians)."""
        r = Basis.from_axis_angle(axis, angle)
        return Basis(r._m @ self._m)




    def scaled(self, s) -> Basis:
        """Return a new basis with columns scaled by s (scalar or 3-tuple)."""
        sv = np.array(s, dtype=np.float32) if not isinstance(s, int | float) else np.array([s, s, s], dtype=np.float32)
        return Basis(self._m * sv[np.newaxis, :])




    def transposed(self) -> Basis:
        """Return transposed basis."""
        return Basis(self._m.T.copy())




    def inverse(self) -> Basis:
        """Return inverse of this basis."""
        return Basis(np.linalg.inv(self._m))




    def get_euler(self) -> tuple[float, float, float]:
        """Extract Euler angles (pitch, yaw, roll) in radians from this basis.

        Uses ZYX intrinsic rotation order (same as engine Quat convention).
        """
        sy = -self._m[2, 0]
        if abs(sy) >= 1.0 - 1e-6:
            yaw = math.asin(max(-1.0, min(1.0, sy)))
            pitch = math.atan2(self._m[0, 1], self._m[0, 2])
            roll = 0.0
        else:
            yaw = math.asin(max(-1.0, min(1.0, sy)))
            pitch = math.atan2(self._m[2, 1], self._m[2, 2])
            roll = math.atan2(self._m[1, 0], self._m[0, 0])
        return (pitch, yaw, roll)




    @classmethod
    def from_euler(cls, euler) -> Basis:
        """Create basis from Euler angles (pitch, yaw, roll) in radians. ZYX order."""
        px = float(euler[0])
        py = float(euler[1])
        pz = float(euler[2])
        cx, sx = math.cos(px), math.sin(px)
        cy, sy = math.cos(py), math.sin(py)
        cz, sz = math.cos(pz), math.sin(pz)
        # R = Rz * Ry * Rx
        return cls(
            np.array(
                [
                    [cy * cz, cz * sx * sy - cx * sz, cx * cz * sy + sx * sz],
                    [cy * sz, cx * cz + sx * sy * sz, cx * sy * sz - cz * sx],
                    [-sy, cy * sx, cx * cy],
                ],
                dtype=np.float32,
            )
        )




    @classmethod
    def from_axis_angle(cls, axis, angle: float) -> Basis:
        """Create rotation basis from axis and angle (radians)."""
        if isinstance(axis, Vec3):
            ax = np.array([axis.x, axis.y, axis.z], dtype=np.float32)
        else:
            ax = np.array(axis, dtype=np.float32)
        ln = np.linalg.norm(ax)
        if ln < 1e-10:
            return cls()
        ax = ax / ln
        c, s = math.cos(angle), math.sin(angle)
        t = 1.0 - c
        x, y, z = ax
        return cls(
            np.array(
                [
                    [t * x * x + c, t * x * y - s * z, t * x * z + s * y],
                    [t * x * y + s * z, t * y * y + c, t * y * z - s * x],
                    [t * x * z - s * y, t * y * z + s * x, t * z * z + c],
                ],
                dtype=np.float32,
            )
        )




    def __eq__(self, other):
        if isinstance(other, Basis):
            return np.allclose(self._m, other._m)
        return NotImplemented




    def __repr__(self):
        return f"Basis({self._m.tolist()})"



# ============================================================================
# Transform3D
# ============================================================================



class Transform3D:
    """3D transform: position + rotation (quat) + scale with cached 4x4 matrix."""

    def __init__(self, position=(0, 0, 0), rotation=None, scale=(1, 1, 1)):
        self._position = np.array(position, dtype=np.float32)
        if rotation is None:
            self._rotation = Quat()
        elif isinstance(rotation, Quat):
            self._rotation = Quat(rotation)
        else:
            self._rotation = Quat(*rotation)
        self._scale = np.array(scale, dtype=np.float32)
        self._matrix: np.ndarray | None = None
        self._dirty = True

    @property
    def position(self) -> np.ndarray:
        return self._position

    @position.setter
    def position(self, value):
        self._position = np.array(value, dtype=np.float32)
        self._dirty = True

    @property
    def rotation(self) -> Quat:
        return self._rotation

    @rotation.setter
    def rotation(self, value: Quat):
        self._rotation = Quat(value)
        self._dirty = True

    @property
    def scale(self) -> np.ndarray:
        return self._scale

    @scale.setter
    def scale(self, value):
        self._scale = np.array(value, dtype=np.float32)
        self._dirty = True

    def _rebuild_matrix(self):
        q = self._rotation
        sx, sy, sz = self._scale
        # Quaternion to rotation matrix entries
        xx = q.x * q.x
        yy = q.y * q.y
        zz = q.z * q.z
        xy = q.x * q.y
        xz = q.x * q.z
        yz = q.y * q.z
        wx = q.w * q.x
        wy = q.w * q.y
        wz = q.w * q.z
        self._matrix = np.array(
            [
                [(1 - 2 * (yy + zz)) * sx, (2 * (xy - wz)) * sy, (2 * (xz + wy)) * sz, self._position[0]],
                [(2 * (xy + wz)) * sx, (1 - 2 * (xx + zz)) * sy, (2 * (yz - wx)) * sz, self._position[1]],
                [(2 * (xz - wy)) * sx, (2 * (yz + wx)) * sy, (1 - 2 * (xx + yy)) * sz, self._position[2]],
                [0, 0, 0, 1],
            ],
            dtype=np.float32,
        )

    @property
    def matrix(self) -> np.ndarray:
        if self._dirty:
            self._rebuild_matrix()
            self._dirty = False
        return self._matrix



    def translated(self, offset) -> Transform3D:
        """Return a new transform with position shifted by offset."""
        return Transform3D(self._position + np.array(offset, dtype=np.float32), self._rotation, self._scale)




    def rotated(self, axis, angle: float) -> Transform3D:
        """Return a new transform with additional rotation (angle in radians)."""
        q = Quat.from_axis_angle(axis if isinstance(axis, Vec3) else Vec3(axis), angle)
        return Transform3D(self._position, q * self._rotation, self._scale)




    def scaled(self, factor) -> Transform3D:
        """Return a new transform with scale multiplied by factor."""
        if isinstance(factor, int | float):
            f = np.array([factor, factor, factor], dtype=np.float32)
        else:
            f = np.array(factor, dtype=np.float32)
        return Transform3D(self._position, self._rotation, self._scale * f)




    def transform_point(self, point) -> np.ndarray:
        """Transform a 3D point through this transform."""
        p = np.array(point, dtype=np.float32)
        h = np.array([p[0], p[1], p[2], 1.0], dtype=np.float32)
        result = self.matrix @ h
        return result[:3]




    def looking_at(self, target, up=(0, 1, 0)) -> Transform3D:
        """Return a new transform oriented to look at target."""
        if not isinstance(target, Vec3):
            target = Vec3(target)
        pos = Vec3(float(self._position[0]), float(self._position[1]), float(self._position[2]))
        direction = target - pos
        if direction.length() < 1e-10:
            return Transform3D(self._position, self._rotation, self._scale)
        q = Quat.look_at(direction, up if isinstance(up, Vec3) else Vec3(up))
        return Transform3D(self._position, q, self._scale)




    def inverse(self) -> Transform3D:
        """Return the inverse transform."""
        inv_q = self._rotation.inverse()
        inv_s = np.where(np.abs(self._scale) > 1e-10, 1.0 / self._scale, np.zeros_like(self._scale))
        # Inverse position: -inv_q * (inv_s * pos)
        sp = Vec3(
            float(self._position[0] * inv_s[0]),
            float(self._position[1] * inv_s[1]),
            float(self._position[2] * inv_s[2]),
        )
        rp = inv_q * sp
        return Transform3D((-rp.x, -rp.y, -rp.z), inv_q, inv_s)


    @property
    def basis(self) -> Basis:
        """Return the 3x3 rotation/scale basis of this transform."""
        return Basis(self.matrix[:3, :3].copy())



    def __mul__(self, other):
        if isinstance(other, Transform3D):
            m = self.matrix @ other.matrix
            return Transform3D._from_matrix(m)
        if isinstance(other, tuple | list | np.ndarray):
            return self.transform_point(other)
        return NotImplemented


    @staticmethod
    def _from_matrix(m: np.ndarray) -> Transform3D:
        """Decompose a 4x4 matrix into Transform3D (position + rotation + scale)."""
        pos = m[:3, 3]
        col0 = m[:3, 0]
        col1 = m[:3, 1]
        col2 = m[:3, 2]
        sx = np.linalg.norm(col0)
        sy = np.linalg.norm(col1)
        sz = np.linalg.norm(col2)
        # Build pure rotation matrix
        rot = np.eye(3, dtype=np.float32)
        if sx > 1e-10:
            rot[:, 0] = col0 / sx
        if sy > 1e-10:
            rot[:, 1] = col1 / sy
        if sz > 1e-10:
            rot[:, 2] = col2 / sz
        # Extract quaternion from rotation matrix
        b = Basis(rot)
        euler = b.get_euler()  # radians
        q = Quat.from_euler(euler[0], euler[1], euler[2])
        return Transform3D(pos, q, (sx, sy, sz))



    def __eq__(self, other):
        if isinstance(other, Transform3D):
            return (
                np.allclose(self._position, other._position)
                and np.allclose(
                    [self._rotation.w, self._rotation.x, self._rotation.y, self._rotation.z],
                    [other._rotation.w, other._rotation.x, other._rotation.y, other._rotation.z],
                    atol=1e-5,
                )
                and np.allclose(self._scale, other._scale)
            )
        return NotImplemented




    def __repr__(self):
        return f"Transform3D(pos={self._position}, rot={self._rotation}, scale={self._scale})"