"""Shared NumPy quaternion operations using the xyzw convention."""

from __future__ import annotations

import numpy as np

from .native_backend import native_kernel


def _mapped(values, operation, width):
    kernel = native_kernel("quaternion_map")
    if kernel is None:
        return None
    array = np.asarray(values, dtype=np.float64)
    return kernel(array.reshape(-1, array.shape[-1]), operation).reshape(*array.shape[:-1], width)


def normalize_quaternion(values: np.ndarray) -> np.ndarray:
    """Normalize quaternion rows and reject an undefined zero rotation."""
    result = _mapped(values, "normalize", 4)
    if result is not None:
        return result
    quaternions = np.asarray(values, dtype=np.float64)
    norms = np.linalg.norm(quaternions, axis=-1, keepdims=True)
    if np.any(norms < 1e-12):
        raise ValueError("a pose contains a zero quaternion")
    return quaternions / norms


def multiply_quaternions(left: np.ndarray, right: np.ndarray) -> np.ndarray:
    """Return the Hamilton product of broadcastable xyzw arrays."""
    first = np.asarray(left, dtype=np.float64)
    second = np.asarray(right, dtype=np.float64)
    kernel = native_kernel("quaternion_multiply")
    if kernel is not None:
        first, second = np.broadcast_arrays(first, second)
        return kernel(first.reshape(-1, 4), second.reshape(-1, 4)).reshape(first.shape)
    first_vector, first_scalar = first[..., :3], first[..., 3:4]
    second_vector, second_scalar = second[..., :3], second[..., 3:4]
    return normalize_quaternion(
        np.concatenate(
            (
                first_scalar * second_vector
                + second_scalar * first_vector
                + np.cross(first_vector, second_vector),
                first_scalar * second_scalar
                - np.sum(first_vector * second_vector, axis=-1, keepdims=True),
            ),
            axis=-1,
        )
    )


def conjugate_quaternion(quaternion: np.ndarray) -> np.ndarray:
    """Return the conjugate of an xyzw quaternion array."""
    native = _mapped(quaternion, "conjugate", 4)
    if native is not None:
        return native
    result = np.asarray(quaternion, dtype=np.float64).copy()
    result[..., :3] *= -1.0
    return result


def rotvec_to_quaternion(rotvec: np.ndarray) -> np.ndarray:
    """Convert rotation vectors to xyzw quaternions."""
    result = _mapped(rotvec, "from_rotvec", 4)
    if result is not None:
        return result
    values = np.asarray(rotvec, dtype=np.float64)
    angle = np.linalg.norm(values, axis=-1)
    scale = np.full_like(angle, 0.5)
    regular = angle >= 1e-12
    scale[regular] = np.sin(0.5 * angle[regular]) / angle[regular]
    return normalize_quaternion(
        np.concatenate((values * scale[..., None], np.cos(0.5 * angle)[..., None]), axis=-1)
    )


def quaternion_to_rotvec(quaternion: np.ndarray) -> np.ndarray:
    """Convert xyzw quaternions to shortest-path rotation vectors."""
    result = _mapped(quaternion, "to_rotvec", 3)
    if result is not None:
        return result
    values = normalize_quaternion(quaternion).copy()
    values[values[..., 3] < 0.0] *= -1.0
    vector_norm = np.linalg.norm(values[..., :3], axis=-1)
    angle = 2.0 * np.arctan2(vector_norm, np.clip(values[..., 3], 0.0, 1.0))
    scale = np.full_like(angle, 2.0)
    regular = vector_norm >= 1e-12
    scale[regular] = angle[regular] / vector_norm[regular]
    return values[..., :3] * scale[..., None]


def euler_xyz_to_quaternion(euler: np.ndarray) -> np.ndarray:
    """Convert ``(N,3)`` XYZ Euler angles to xyzw quaternions."""
    values = np.asarray(euler, dtype=np.float64)
    if values.ndim != 2 or values.shape[1] != 3:
        raise ValueError("XYZ Euler angles must have shape [N, 3]")
    result = _mapped(values, "from_euler", 4)
    if result is not None:
        return result
    half = 0.5 * values
    sin_roll, sin_pitch, sin_yaw = np.sin(half).T
    cos_roll, cos_pitch, cos_yaw = np.cos(half).T
    result = np.empty((len(values), 4), dtype=np.float64)
    result[:, 0] = sin_roll * cos_pitch * cos_yaw - cos_roll * sin_pitch * sin_yaw
    result[:, 1] = cos_roll * sin_pitch * cos_yaw + sin_roll * cos_pitch * sin_yaw
    result[:, 2] = cos_roll * cos_pitch * sin_yaw - sin_roll * sin_pitch * cos_yaw
    result[:, 3] = cos_roll * cos_pitch * cos_yaw + sin_roll * sin_pitch * sin_yaw
    return normalize_quaternion(result)


def quaternion_to_euler_xyz(quaternion: np.ndarray) -> np.ndarray:
    """Convert ``(N,4)`` xyzw quaternions to XYZ Euler angles."""
    values = normalize_quaternion(quaternion)
    if values.ndim != 2 or values.shape[1] != 4:
        raise ValueError("xyzw quaternions must have shape [N, 4]")
    result = _mapped(values, "to_euler", 3)
    if result is not None:
        return result
    x, y, z, w = values.T
    return np.stack(
        (
            np.arctan2(2.0 * (w * x + y * z), 1.0 - 2.0 * (x * x + y * y)),
            np.arcsin(np.clip(2.0 * (w * y - z * x), -1.0, 1.0)),
            np.arctan2(2.0 * (w * z + x * y), 1.0 - 2.0 * (y * y + z * z)),
        ),
        axis=1,
    )