"""Rotation utilities for observable computations.

Provides helper functions for working with scipy Rotation objects and quaternion arrays.
"""

from __future__ import annotations

from typing import TYPE_CHECKING

import numpy as np
from scipy.spatial.transform import Rotation

if TYPE_CHECKING:
    from collections.abc import Sequence

    from numpy.typing import NDArray


def rotation_to_quat_array(
    rotation_input: NDArray[np.floating] | Sequence | Rotation,
) -> NDArray[np.floating]:
    """Convert rotation input to quaternion array, preserving shape for broadcasting.

    Handles both scipy Rotation objects and quaternion arrays, returning a consistent
    numpy array format suitable for broadcasting operations.

    Args:
        rotation_input: Either a scipy Rotation object, or a quaternion array with
            shape (4,) for single rotation or (..., 4) for batch. Uses scipy
            convention [x, y, z, w] (scalar-last).

    Returns:
        Quaternion array with shape (4,) for single rotation or (..., 4) for batch.
    """
    if isinstance(rotation_input, Rotation):
        rot: Rotation = rotation_input
        quat = rot.as_quat()
        # Single rotation gives (4,), batch gives (n, 4)
        if rot.single:
            return quat.reshape(4)
        return quat
    return np.asarray(rotation_input)
