Math Classes

class Vec2(VectorBase["Vec2"]):
    """A simple 2D vector class for graphics, using numpy for efficient operations.

    Attributes:
        x (float): The x-coordinate of the vector.
        y (float): The y-coordinate of the vector.
    """

    DIMENSION = 2
    COMPONENT_NAMES = ("x", "y")
    DEFAULT_VALUES = (0.0, 0.0)

    __slots__ = ["_data"]

    def cross(self, rhs: "Vec2") -> float:
        """Cross product of two vectors a x b (2D version returns scalar).

        Args:
            rhs (Vec2): The right-hand side vector to cross product with.

        Returns:
            float: 2D cross product (perpendicular dot product).
        """
        return self._data[0] * rhs._data[1] - self._data[1] * rhs._data[0]

    def reflected(self, n: "Vec2") -> "Vec2":
        """Return a new vector reflected about a normal.

        Args:
            n (Vec2): The normal to reflect about.

        Returns:
            Vec2: A new vector that is the result of reflecting this vector about the normal.
        """
        d = self.dot(n)
        # I - 2.0 * dot(N, I) * N
        return Vec2(
            self._data[0] - 2.0 * d * n._data[0], self._data[1] - 2.0 * d * n._data[1]
        )

    def outer(self, rhs: "Vec2") -> "Mat2":
        """Outer product of two vectors a x b.

        Args:
            rhs (Vec2): The right-hand side vector to outer product with.

        Returns:
            Mat2: A new 2x2 matrix that is the result of the outer product.
        """
        from .mat2 import Mat2

        result = Mat2()
        result._data = np.outer(self._data, rhs._data).astype(np.float32)
        return result

    def __matmul__(self, rhs: "Mat2") -> "Vec2":
        """Vec2 @ Mat2 matrix multiplication.

        Args:
            rhs (Mat2): The matrix to multiply by.

        Returns:
            Vec2: A new vector that is the result of multiplying this vector by the matrix.
        """
        return Vec2(
            self._data[0] * rhs._data[0, 0] + self._data[1] * rhs._data[1, 0],
            self._data[0] * rhs._data[0, 1] + self._data[1] * rhs._data[1, 1],
        )

    def set(self, *args: float) -> None:
        """Set the x,y values of the vector.

        Args:
            *args: Component values (x, y).

        Raises:
            ValueError: If wrong number of arguments or they are not floats.
        """
        if len(args) != 2:
            raise ValueError(f"Vec2.set requires 2 arguments, got {len(args)}")
        try:
            self._data[0] = float(args[0])
            self._data[1] = float(args[1])
        except ValueError:
            raise ValueError(f"Vec2.set {args=} all need to be float")

class Vec3(VectorBase["Vec3"]):
    """A simple 3D vector class for 3D graphics, using numpy for efficient operations.

    Attributes:
        x (float): The x-coordinate of the vector.
        y (float): The y-coordinate of the vector.
        z (float): The z-coordinate of the vector.
    """

    DIMENSION = 3
    COMPONENT_NAMES = ("x", "y", "z")
    DEFAULT_VALUES = (0.0, 0.0, 0.0)

    __slots__ = ["_data"]

    def cross(self, rhs: "Vec3") -> "Vec3":
        """Cross product of two vectors a x b.

        Args:
            rhs (Vec3): The right-hand side vector to cross product with.

        Returns:
            Vec3: A new vector that is the result of the cross product.
        """
        result = Vec3()
        result._data = np.cross(self._data, rhs._data)
        return result

    def reflected(self, n: "Vec3") -> "Vec3":
        """Return a new vector reflected about a normal.

        Args:
            n (Vec3): The normal to reflect about.

        Returns:
            Vec3: A new vector that is the result of reflecting this vector about the normal.
        """
        d = self.dot(n)
        # I - 2.0 * dot(N, I) * N
        result = Vec3()
        result._data = self._data - 2.0 * d * n._data
        return result

    def outer(self, rhs: "Vec3") -> "Mat3":
        """Outer product of two vectors a x b.

        Args:
            rhs (Vec3): The right-hand side vector to outer product with.

        Returns:
            Mat3: A new 3x3 matrix that is the result of the outer product.
        """
        from .mat3 import Mat3

        result = Mat3()
        result._data = np.outer(self._data, rhs._data).astype(np.float32)
        return result

    def __matmul__(self, rhs: "Mat3") -> "Vec3":
        """Vec3 @ Mat3 matrix multiplication.

        Args:
            rhs (Mat3): The matrix to multiply by.

        Returns:
            Vec3: A new vector that is the result of multiplying this vector by the matrix.
        """
        result = Vec3()
        result._data = rhs._data.T @ self._data  # More efficient
        return result

    def set(self, *args: float) -> None:
        """Set the x,y,z values of the vector.

        Args:
            *args: Component values (x, y, z).

        Raises:
            ValueError: If wrong number of arguments or they are not floats.
        """
        if len(args) != 3:
            raise ValueError(f"Vec3.set requires 3 arguments, got {len(args)}")
        try:
            self._data[0] = float(args[0])
            self._data[1] = float(args[1])
            self._data[2] = float(args[2])
        except ValueError:
            raise ValueError(f"Vec3.set {args=} all need to be float")

class Vec4(VectorBase["Vec4"]):
    """A simple 4D vector class for graphics, using numpy for efficient operations.

    Attributes:
        x (float): The x-coordinate of the vector.
        y (float): The y-coordinate of the vector.
        z (float): The z-coordinate of the vector.
        w (float): The w-coordinate of the vector.
    """

    DIMENSION = 4
    COMPONENT_NAMES = ("x", "y", "z", "w")
    DEFAULT_VALUES = (0.0, 0.0, 0.0, 1.0)

    __slots__ = ["_data"]

    def cross(self, rhs: "Vec4") -> "Vec4":
        """Cross product of two vectors a x b (4D version uses first 3 components).

        Args:
            rhs (Vec4): The right-hand side vector to cross product with.

        Returns:
            Vec4: A new vector that is the result of the cross product.
        """
        result = Vec4()
        # Cross product only makes sense for 3D vectors, use first 3 components
        result._data[:3] = np.cross(self._data[:3], rhs._data[:3])
        result._data[3] = 0.0
        return result

    def reflected(self, n: "Vec4") -> "Vec4":
        """Return a new vector reflected about a normal.

        Args:
            n (Vec4): The normal to reflect about.

        Returns:
            Vec4: A new vector that is the result of reflecting this vector about the normal.
        """
        d = self.dot(n)
        # I - 2.0 * dot(N, I) * N
        result = Vec4()
        result._data = self._data - 2.0 * d * n._data
        return result

    def outer(self, rhs: "Vec4") -> "Mat4":
        """Outer product of two vectors a x b.

        Args:
            rhs (Vec4): The right-hand side vector to outer product with.

        Returns:
            Mat4: A new 4x4 matrix that is the result of the outer product.
        """
        from .mat4 import Mat4

        result = Mat4()
        result._data = np.outer(self._data, rhs._data).astype(np.float32)
        return result

    def __matmul__(self, rhs: "Mat4") -> "Vec4":
        """Vec4 @ Mat4 matrix multiplication.

        Args:
            rhs (Mat4): The matrix to multiply by.

        Returns:
            Vec4: A new vector that is the result of multiplying this vector by the matrix.
        """
        return Vec4(*(self._data @ rhs._data))

    def set(self, *args: float) -> None:
        """Set the x,y,z,w values of the vector.

        Args:
            *args: Component values (x, y, z, w). w defaults to 1.0 if not provided.

        Raises:
            ValueError: If wrong number of arguments or they are not floats.
        """
        if len(args) == 3:
            # Allow (x, y, z) with default w=1.0 for backward compatibility
            args = args + (1.0,)
        elif len(args) != 4:
            raise ValueError(f"Vec4.set requires 3 or 4 arguments, got {len(args)}")

        try:
            for i in range(4):
                self._data[i] = float(args[i])
        except ValueError:
            raise ValueError(f"Vec4.set {args=} all need to be float")

class Mat2(MatrixBase):
    """A 2x2 matrix for 2D transforms."""

    SIZE = 2

    def _vec_type(self) -> type:
        from .vec2 import Vec2

        return Vec2

class Mat3(MatrixBase):
    """A 3x3 matrix for basic affine transforms."""

    SIZE = 3

    def _vec_type(self) -> type:
        from .vec3 import Vec3

        return Vec3

    @classmethod
    def scale(cls, x: float, y: float, z: float) -> "Mat3":
        """Return a scale matrix with the diagonal set to (x, y, z)."""
        a = cls()
        a._data[0, 0] = x
        a._data[1, 1] = y
        a._data[2, 2] = z
        return a

    @classmethod
    def rotate_x(cls, angle: float) -> "Mat3":
        """Return a rotation matrix around the X axis by angle degrees."""
        a = cls()
        beta = math.radians(angle)
        sr = math.sin(beta)
        cr = math.cos(beta)
        a._data[1, 1] = cr
        a._data[1, 2] = sr
        a._data[2, 1] = -sr
        a._data[2, 2] = cr
        return a

    @classmethod
    def rotate_y(cls, angle: float) -> "Mat3":
        """Return a rotation matrix around the Y axis by angle degrees."""
        a = cls()
        beta = math.radians(angle)
        sr = math.sin(beta)
        cr = math.cos(beta)
        a._data[0, 0] = cr
        a._data[0, 2] = -sr
        a._data[2, 0] = sr
        a._data[2, 2] = cr
        return a

    @classmethod
    def rotate_z(cls, angle: float) -> "Mat3":
        """Return a rotation matrix around the Z axis by angle degrees."""
        a = cls()
        beta = math.radians(angle)
        sr = math.sin(beta)
        cr = math.cos(beta)
        a._data[0, 0] = cr
        a._data[0, 1] = sr
        a._data[1, 0] = -sr
        a._data[1, 1] = cr
        return a

    @classmethod
    def from_mat4(cls, mat4: "Mat4") -> "Mat3":
        """Return the upper-left 3x3 of a Mat4."""
        result = cls()
        result._data = mat4._data[:3, :3].copy()
        return result

class Mat4(MatrixBase):
    """A 4x4 matrix for 3D affine and projective transforms."""

    SIZE = 4

    def _vec_type(self) -> type:
        from .vec4 import Vec4

        return Vec4

    @classmethod
    def scale(cls, x: float, y: float, z: float) -> "Mat4":
        """Return a scale matrix with the diagonal set to (x, y, z, 1)."""
        a = cls()
        a._data[0, 0] = x
        a._data[1, 1] = y
        a._data[2, 2] = z
        return a

    @classmethod
    def translate(cls, x: float, y: float, z: float) -> "Mat4":
        """Return a translation matrix."""
        a = cls()
        a._data[3, 0] = x
        a._data[3, 1] = y
        a._data[3, 2] = z
        return a

    @classmethod
    def rotate_x(cls, angle: float) -> "Mat4":
        """Return a rotation matrix around the X axis by angle degrees."""
        a = cls()
        beta = math.radians(angle)
        sr = math.sin(beta)
        cr = math.cos(beta)
        a._data[1, 1] = cr
        a._data[1, 2] = sr
        a._data[2, 1] = -sr
        a._data[2, 2] = cr
        return a

    @classmethod
    def rotate_y(cls, angle: float) -> "Mat4":
        """Return a rotation matrix around the Y axis by angle degrees."""
        a = cls()
        beta = math.radians(angle)
        sr = math.sin(beta)
        cr = math.cos(beta)
        a._data[0, 0] = cr
        a._data[0, 2] = -sr
        a._data[2, 0] = sr
        a._data[2, 2] = cr
        return a

    @classmethod
    def rotate_z(cls, angle: float) -> "Mat4":
        """Return a rotation matrix around the Z axis by angle degrees."""
        a = cls()
        beta = math.radians(angle)
        sr = math.sin(beta)
        cr = math.cos(beta)
        a._data[0, 0] = cr
        a._data[0, 1] = sr
        a._data[1, 0] = -sr
        a._data[1, 1] = cr
        return a

    @classmethod
    def from_mat3(cls, mat3: "Mat3") -> "Mat4":
        """Return a Mat4 with the given Mat3 as its upper-left block."""
        result = cls()
        result._data[:3, :3] = mat3._data
        return result

class Quaternion:
    """A quaternion for representing and composing 3D rotations.

    Attributes:
        s (float): The scalar part of the quaternion.
        x (float): The x-coordinate of the vector part of the quaternion.
        y (float): The y-coordinate of the vector part of the quaternion.
        z (float): The z-coordinate of the vector part of the quaternion.
    """

    __slots__ = ("_data",)  # Store as [s, x, y, z]

    def __init__(
        self, s: float = 1.0, x: float = 0.0, y: float = 0.0, z: float = 0.0
    ) -> None:
        """Initializes a new instance of the Quaternion class.

        Args:
            s (float): The scalar part of the quaternion.
            x (float): The x-coordinate of the vector part of the quaternion.
            y (float): The y-coordinate of the vector part of the quaternion.
            z (float): The z-coordinate of the vector part of the quaternion.
        """
        self._data = np.array(
            [float(s), float(x), float(y), float(z)], dtype=np.float32
        )

    @classmethod
    def from_mat4(cls, mat: "Mat4") -> "Quaternion":
        """Creates a new Quaternion from a Mat4 rotation matrix.

        Args:
            mat (Mat4): The rotation matrix to convert.

        Returns:
            Quaternion: A new Quaternion representing the rotation matrix.
        """
        matrix = mat.to_list()
        T = 1.0 + matrix[0] + matrix[5] + matrix[10]
        if T > 0.00000001:  # to avoid large distortions!
            scale = math.sqrt(T) * 2.0
            x = (matrix[6] - matrix[9]) / scale
            y = (matrix[8] - matrix[2]) / scale
            z = (matrix[1] - matrix[4]) / scale
            s = 0.25 * scale
        elif matrix[0] > matrix[5] and matrix[0] > matrix[10]:
            scale = math.sqrt(1.0 + matrix[0] - matrix[5] - matrix[10]) * 2.0
            x = 0.25 * scale
            y = (matrix[4] + matrix[1]) / scale
            z = (matrix[2] + matrix[8]) / scale
            s = (matrix[6] - matrix[9]) / scale
        elif matrix[5] > matrix[10]:
            scale = math.sqrt(1.0 + matrix[5] - matrix[0] - matrix[10]) * 2.0
            x = (matrix[4] + matrix[1]) / scale
            y = 0.25 * scale
            z = (matrix[9] + matrix[6]) / scale
            s = (matrix[8] - matrix[2]) / scale
        else:
            scale = math.sqrt(1.0 + matrix[10] - matrix[0] - matrix[5]) * 2.0
            x = (matrix[8] + matrix[2]) / scale
            y = (matrix[9] + matrix[6]) / scale
            z = 0.25 * scale
            s = (matrix[1] - matrix[4]) / scale

        return cls(s, x, y, z)

    @classmethod
    def from_axis_angle(cls, axis: "Vec3", angle: float) -> "Quaternion":
        """Creates a new Quaternion from an axis and angle.

        Args:
            axis (Vec3): The axis of rotation.
            angle (float): The angle of rotation in degrees.

        Returns:
            Quaternion: A new Quaternion representing the rotation.
        """
        angle_rad = math.radians(angle)
        half_angle = angle_rad * 0.5
        s = math.cos(half_angle)
        sin_half_angle = math.sin(half_angle)
        x = axis.x * sin_half_angle
        y = axis.y * sin_half_angle
        z = axis.z * sin_half_angle
        return cls(s, x, y, z)

    def __add__(self, rhs: "Quaternion") -> "Quaternion":
        """Quaternion addition a+b, component-wise."""
        result = Quaternion()
        result._data = self._data + rhs._data
        return result

    def __sub__(self, rhs: "Quaternion") -> "Quaternion":
        """Quaternion subtraction a-b, component-wise."""
        result = Quaternion()
        result._data = self._data - rhs._data
        return result

    def __matmul__(self, rhs: "Quaternion") -> "Quaternion":
        """Quaternion product (Hamilton), returning a new quaternion."""
        if not isinstance(rhs, Quaternion):
            raise TypeError("@ requires a Quaternion")
        s1, x1, y1, z1 = self._data
        s2, x2, y2, z2 = rhs._data
        return Quaternion(
            s1 * s2 - x1 * x2 - y1 * y2 - z1 * z2,
            s1 * x2 + x1 * s2 + y1 * z2 - z1 * y2,
            s1 * y2 - x1 * z2 + y1 * s2 + z1 * x2,
            s1 * z2 + x1 * y2 - y1 * x2 + z1 * s2,
        )

    def __mul__(self, rhs: "float | int | Vec3") -> "Quaternion | Vec3":
        """Scalar scale or Vec3 rotation. Quaternion product uses @."""
        if isinstance(rhs, Quaternion):
            raise TypeError("use q1 @ q2 for the quaternion product")
        if isinstance(rhs, (int, float)):
            result = Quaternion()
            result._data = self._data * np.float32(rhs)
            return result
        if isinstance(rhs, Vec3):
            # Quaternion-vector multiplication (rotate vector by quaternion)
            qw = self.s
            qx = self.x
            qy = self.y
            qz = self.z

            vx = rhs.x
            vy = rhs.y
            vz = rhs.z

            # pq (quaternion * pure quaternion from vector)
            pw = -qx * vx - qy * vy - qz * vz
            px = qw * vx + qy * vz - qz * vy
            py = qw * vy - qx * vz + qz * vx
            pz = qw * vz + qx * vy - qy * vx

            # pqp* (result * conjugate of quaternion)
            return Vec3(
                -pw * qx + px * qw - py * qz + pz * qy,
                -pw * qy + px * qz + py * qw - pz * qx,
                -pw * qz - px * qy + py * qx + pz * qw,
            )
        raise TypeError(f"cannot multiply Quaternion by {type(rhs)}")

    def __rmul__(self, rhs: float) -> "Quaternion":
        """Scalar scale (right operand)."""
        if isinstance(rhs, (int, float)):
            return self * rhs
        raise TypeError(f"cannot multiply {type(rhs)} by Quaternion")

    def __neg__(self) -> "Quaternion":
        """Return a new quaternion with every component negated."""
        result = Quaternion()
        result._data = -self._data
        return result

    def __truediv__(self, rhs: float | int) -> "Quaternion":
        """Scalar division, returning a new quaternion.

        Raises:
            ZeroDivisionError: If rhs is zero.
            TypeError: If rhs is not a scalar.
        """
        if isinstance(rhs, (int, float)):
            if rhs == 0:
                raise ZeroDivisionError("division by zero")
            result = Quaternion()
            result._data = self._data / np.float32(rhs)
            return result
        raise TypeError(f"cannot divide Quaternion by {type(rhs)}")

    def __getitem__(self, index: int) -> float:
        """Return the component at index (0=s, 1=x, 2=y, 3=z).

        Raises:
            IndexError: If the index is out of range.
        """
        if index < 0 or index >= 4:
            raise IndexError("Index out of range. Valid indices are 0, 1, 2, 3.")
        return float(self._data[index])

    def normalized(self) -> "Quaternion":
        """Return a new unit-length quaternion.

        Raises:
            ZeroDivisionError: If the quaternion has zero length.
        """
        length = self.length()
        if math.isclose(length, 0.0):
            raise ZeroDivisionError("Quaternion.normalized: length is zero")
        result = Quaternion()
        result._data = self._data / np.float32(length)
        return result

    def length(self) -> float:
        """Return the length/magnitude of the quaternion."""
        return float(np.linalg.norm(self._data))

    def length_squared(self) -> float:
        """Return the squared magnitude."""
        return float(np.dot(self._data, self._data))

    def conjugate(self) -> "Quaternion":
        """Return the conjugate of the quaternion (s, -x, -y, -z)."""
        result = Quaternion()
        result._data = self._data.copy()
        result._data[1:] *= -1  # Negate x, y, z components
        return result

    def inverse(self) -> "Quaternion":
        """Return the multiplicative inverse (conjugate / |q|^2)."""
        lsq = self.length_squared()
        if math.isclose(lsq, 0.0):
            raise ZeroDivisionError("Quaternion.inverse: zero quaternion")
        result = self.conjugate()
        result._data = result._data / np.float32(lsq)
        return result

    def dot(self, rhs: "Quaternion") -> float:
        """Dot product of two quaternions."""
        return float(np.dot(self._data, rhs._data))

    def slerp(self, rhs: "Quaternion", t: float) -> "Quaternion":
        """Spherical linear interpolation from self to rhs at t in [0, 1]."""
        dot = float(np.dot(self._data, rhs._data))
        rhs_data = rhs._data.copy()
        if dot < 0.0:
            dot = -dot
            rhs_data = -rhs_data
        if dot > 0.9995:
            data = self._data + np.float32(t) * (rhs_data - self._data)
            data = data / np.linalg.norm(data)
        else:
            theta0 = math.acos(max(-1.0, min(1.0, dot)))
            theta = theta0 * t
            s0 = math.cos(theta) - dot * math.sin(theta) / math.sin(theta0)
            s1 = math.sin(theta) / math.sin(theta0)
            data = np.float32(s0) * self._data + np.float32(s1) * rhs_data
        result = Quaternion()
        result._data = data.astype(np.float32)
        return result

    def to_mat4(self) -> Mat4:
        """Return the equivalent rotation matrix (row-vector convention)."""
        s, x, y, z = (float(v) for v in self._data)
        m = Mat4()
        m._data[0, 0] = 1.0 - 2.0 * (y * y + z * z)
        m._data[0, 1] = 2.0 * (x * y + s * z)
        m._data[0, 2] = 2.0 * (x * z - s * y)
        m._data[1, 0] = 2.0 * (x * y - s * z)
        m._data[1, 1] = 1.0 - 2.0 * (x * x + z * z)
        m._data[1, 2] = 2.0 * (y * z + s * x)
        m._data[2, 0] = 2.0 * (x * z + s * y)
        m._data[2, 1] = 2.0 * (y * z - s * x)
        m._data[2, 2] = 1.0 - 2.0 * (x * x + y * y)
        return m

    def set(self, s: float, x: float, y: float, z: float) -> None:
        """Set all four components."""
        self._data[:] = (float(s), float(x), float(y), float(z))

    def copy(self) -> "Quaternion":
        """Return a new quaternion with the same values."""
        return Quaternion(*self._data)

    def to_numpy(self) -> np.ndarray:
        """Return the quaternion as a numpy array [s, x, y, z]."""
        return self._data.copy()

    def to_list(self) -> list[float]:
        """Return the quaternion as a list [s, x, y, z]."""
        return self._data.tolist()

    def to_tuple(self) -> tuple[float, float, float, float]:
        """Return (s, x, y, z) as plain floats."""
        return tuple(float(v) for v in self._data)

    @classmethod
    def from_list(cls, lst: list[float]) -> "Quaternion":
        """Create from [s, x, y, z]."""
        if len(lst) != 4:
            raise ValueError("Quaternion.from_list requires 4 values")
        return cls(*lst)

    @classmethod
    def from_numpy(cls, arr: np.ndarray) -> "Quaternion":
        """Create from an array [s, x, y, z]."""
        arr = np.asarray(arr, dtype=np.float32)
        if arr.shape != (4,):
            raise ValueError("Quaternion.from_numpy requires shape (4,)")
        return cls(*arr)

    def __eq__(self, rhs: object) -> bool:
        """Quaternion comparison a==b using numpy.allclose."""
        if not isinstance(rhs, Quaternion):
            return NotImplemented
        return bool(np.allclose(self._data, rhs._data, rtol=1e-5, atol=1e-6))

    def __ne__(self, rhs: object) -> bool:
        """Quaternion comparison a!=b using numpy.allclose."""
        result = self.__eq__(rhs)
        return result if result is NotImplemented else not result

    def __hash__(self) -> int:
        """Compute hash for use in sets and dictionaries."""
        from .util import hash_combine

        seed = 0
        for v in self._data:
            seed = hash_combine(seed, hash(float(np.float32(v))))
        return seed

    def __len__(self) -> int:
        """Return the number of components (always 4)."""
        return 4

    def __iter__(self) -> Generator[float, None, None]:
        """Yield the components in order (s, x, y, z)."""
        return iter(self._data.tolist())

    def __repr__(self) -> str:
        """Eval-able representation, e.g. Quaternion(1.0, 0.0, 0.0, 0.0)."""
        args = ", ".join(repr(float(v)) for v in self._data)
        return f"Quaternion({args})"

    def __str__(self) -> str:
        """Pretty representation, e.g. Quaternion(1.0, [0.0, 0.0, 0.0])."""
        s, x, y, z = (float(v) for v in self._data)
        return f"Quaternion({s}, [{x}, {y}, {z}])"