"""Small numpy quaternion/rotation helpers (quaternions are wxyz, as in Isaac Lab)."""

from __future__ import annotations

import math

import numpy as np


def quat_normalize(q: np.ndarray) -> np.ndarray:
    n = np.linalg.norm(q)
    return q / n if n > 1e-8 else np.array([1.0, 0.0, 0.0, 0.0], dtype=q.dtype)


def quat_mul(a: np.ndarray, b: np.ndarray) -> np.ndarray:
    w1, x1, y1, z1 = a
    w2, x2, y2, z2 = b
    return np.array(
        [
            w1 * w2 - x1 * x2 - y1 * y2 - z1 * z2,
            w1 * x2 + x1 * w2 + y1 * z2 - z1 * y2,
            w1 * y2 - x1 * z2 + y1 * w2 + z1 * x2,
            w1 * z2 + x1 * y2 - y1 * x2 + z1 * w2,
        ],
        dtype=np.float64,
    )


def quat_conj(q: np.ndarray) -> np.ndarray:
    return np.array([q[0], -q[1], -q[2], -q[3]], dtype=np.float64)


def quat_to_matrix(q: np.ndarray) -> np.ndarray:
    w, x, y, z = quat_normalize(np.asarray(q, dtype=np.float64))
    return np.array(
        [
            [1 - 2 * (y * y + z * z), 2 * (x * y - z * w), 2 * (x * z + y * w)],
            [2 * (x * y + z * w), 1 - 2 * (x * x + z * z), 2 * (y * z - x * w)],
            [2 * (x * z - y * w), 2 * (y * z + x * w), 1 - 2 * (x * x + y * y)],
        ],
        dtype=np.float64,
    )


def quat_apply(q: np.ndarray, v: np.ndarray) -> np.ndarray:
    return quat_to_matrix(q) @ np.asarray(v, dtype=np.float64)


def quat_apply_inverse(q: np.ndarray, v: np.ndarray) -> np.ndarray:
    return quat_to_matrix(q).T @ np.asarray(v, dtype=np.float64)


def quat_from_euler_xyz(roll: float, pitch: float, yaw: float) -> np.ndarray:
    """Isaac Lab ``quat_from_euler_xyz`` convention (intrinsic x-y-z), wxyz."""
    cr, sr = math.cos(0.5 * roll), math.sin(0.5 * roll)
    cp, sp = math.cos(0.5 * pitch), math.sin(0.5 * pitch)
    cy, sy = math.cos(0.5 * yaw), math.sin(0.5 * yaw)
    return np.array(
        [
            cr * cp * cy + sr * sp * sy,
            sr * cp * cy - cr * sp * sy,
            cr * sp * cy + sr * cp * sy,
            cr * cp * sy - sr * sp * cy,
        ],
        dtype=np.float64,
    )


def yaw_of(q: np.ndarray) -> float:
    w, x, y, z = q
    return math.atan2(2.0 * (w * z + x * y), 1.0 - 2.0 * (y * y + z * z))


def rot6d_interleaved(q: np.ndarray) -> np.ndarray:
    """``matrix_from_quat(q)[..., :2].reshape(6)`` as used by the ``motion_root_ori`` observation.

    Row-major reshape of the first two columns: ``[m00, m01, m10, m11, m20, m21]``.
    """
    return quat_to_matrix(q)[:, :2].reshape(6)


def quat_slerp_sign_aligned(a: np.ndarray, b: np.ndarray, alpha: float) -> np.ndarray:
    """Linear blend with sign alignment then renormalisation (training blend convention)."""
    b = b if float(np.dot(a, b)) >= 0.0 else -b
    return quat_normalize(a + alpha * (b - a))