# SPDX-License-Identifier: Apache-2.0
"""C10 geometry (sweep and ego-motion parts): rigid transforms in float64, Autoware's float32 sweep transform, and
the yaw / quaternion conventions of the LiDAR detectors.

Naming: ``T_a_from_b`` maps points of frame ``b`` into frame ``a`` (4x4, metres, row-major, translation in the last
column), so ``p_a = T_a_from_b @ p_b`` and ``T_a_from_c = T_a_from_b @ T_b_from_c``.

Two precisions on purpose:

- **float64** for composing poses (``compose``, ``invert_rigid``): what tf2 does, and what every port uses unless its
  Autoware node does otherwise (S:streampetr:386; S:pointpainting:212-215; S:bevformer:437).
- **float32 as Autoware's lidar_centerpoint does it**: the node casts each TF to ``Eigen::Affine3f``, inverts and
  multiplies in float32 (``pointcloud_densification.cpp:47-51,82-86``; ``voxel_generator.cpp:63``) and the CUDA
  kernel ``generateSweepPoints_kernel`` applies the float32 affine as ``m0*x + m4*y + m8*z + m12`` (column-major,
  ``preprocess_kernel.cu:58-90``). :func:`invert_affine_f32` and :func:`transform_points_f32` reproduce that
  arithmetic in numpy float32 without fused multiply-adds, so they agree with the GPU to the last bit except where
  nvcc contracts a multiply-add (an FMA rounds once): a sub-micrometre difference, irrelevant to 0.32 m pillars.

Yaw conventions (S:centerpoint:305-307): the deployed CenterPoint ONNX regresses the "mmdet3d" yaw; Autoware
publishes ``yaw_ros = -yaw_net - pi/2`` counter-clockwise from +x of ``base_link`` (``ros_utils.cpp:55-60``).

numpy only; no side effects on import.
"""
from __future__ import annotations

import math
from typing import Any, Sequence

import numpy as np

from .io import parse_transform

__all__ = [
    "as_transform",
    "compose",
    "invert_rigid",
    "invert_affine_f32",
    "compose_f32",
    "transform_points",
    "transform_points_f32",
    "rotation_from_rpy",
    "transform_from_xyz_rpy",
    "quaternion_wxyz_from_yaw",
    "yaw_from_quaternion_wxyz",
    "yaw_from_rotation",
    "wrap_angle",
    "mmdet_yaw_to_ros",
    "ros_yaw_to_mmdet",
]


# ------------------------------------------------------------------------------------------------- transforms

def as_transform(obj: Any, *, field: str = "transform") -> np.ndarray:
    """Any spelling :func:`ttaw.io.parse_transform` accepts (4x4 / 3x4 array, ``{matrix}``,
    ``{translation, rotation_wxyz | rotation_xyzw | rotation}``, ``{x, y, z, roll, pitch, yaw}``) -> a validated
    (4, 4) float64 rigid transform. A malformed value raises :class:`ttaw.io.InputError`."""
    return parse_transform(obj, field=field)


def compose(*transforms: np.ndarray) -> np.ndarray:
    """``T_0 @ T_1 @ ... @ T_n`` in float64 (``compose(T_a_from_b, T_b_from_c) == T_a_from_c``)."""
    if not transforms:
        return np.eye(4)
    out = np.asarray(transforms[0], dtype=np.float64)
    for t in transforms[1:]:
        out = out @ np.asarray(t, dtype=np.float64)
    return out


def invert_rigid(T: np.ndarray) -> np.ndarray:
    """Exact inverse of a rigid transform in float64: ``[R^T, -R^T t]``."""
    T = np.asarray(T, dtype=np.float64)
    out = np.eye(4)
    r = T[:3, :3]
    out[:3, :3] = r.T
    out[:3, 3] = -(r.T @ T[:3, 3])
    return out


def invert_affine_f32(T: np.ndarray) -> np.ndarray:
    """``Eigen::Affine3f::inverse()`` (TransformTraits ``Affine``) in float32: the 3x3 linear part inverted through
    its cofactors and the column-0 determinant (Eigen's ``compute_inverse<..., 3>``), translation ``-L^-1 t``,
    every operation rounded to float32. Used by the Autoware densification emulation
    (``pointcloud_densification.cpp:84``)."""
    m = np.asarray(T, dtype=np.float32)
    a = m[:3, :3]
    f = np.float32

    def cofactor(i: int, j: int) -> np.float32:  # Eigen cofactor_3x3<i, j>
        i1, i2, j1, j2 = (i + 1) % 3, (i + 2) % 3, (j + 1) % 3, (j + 2) % 3
        return f(f(a[i1, j1] * a[i2, j2]) - f(a[i1, j2] * a[i2, j1]))

    det = f(f(cofactor(0, 0) * a[0, 0]) + f(cofactor(1, 0) * a[1, 0])) + f(cofactor(2, 0) * a[2, 0])
    if det == 0 or not np.isfinite(det):
        raise ValueError("singular affine transform")
    inv_det = f(f(1.0) / det)
    inv = np.empty((3, 3), np.float32)
    for i in range(3):
        for j in range(3):
            inv[i, j] = f(cofactor(j, i) * inv_det)   # adjugate / det
    out = np.eye(4, dtype=np.float32)
    out[:3, :3] = inv
    t = m[:3, 3]
    for i in range(3):
        out[i, 3] = -(f(f(inv[i, 0] * t[0]) + f(inv[i, 1] * t[1])) + f(inv[i, 2] * t[2]))
    return out


def compose_f32(A: np.ndarray, B: np.ndarray) -> np.ndarray:
    """``A @ B`` of two float32 affine matrices with float32 rounding after every product and sum, summed in index
    order (Eigen's lazy product of two ``Affine3f``)."""
    a = np.asarray(A, dtype=np.float32)
    b = np.asarray(B, dtype=np.float32)
    out = np.zeros((4, 4), np.float32)
    for k in range(4):
        out = (out + a[:, k:k + 1] * b[k:k + 1, :]).astype(np.float32)
    return out


def transform_points(points: np.ndarray, T: np.ndarray) -> np.ndarray:
    """``T @ [x, y, z, 1]`` for (N, >=3) points in float64 (extra columns are passed through unchanged)."""
    p = np.asarray(points, dtype=np.float64)
    T = np.asarray(T, dtype=np.float64)
    out = p.copy()
    out[:, :3] = p[:, :3] @ T[:3, :3].T + T[:3, 3]
    return out


def transform_points_f32(xyz: np.ndarray, T: np.ndarray) -> np.ndarray:
    """Autoware ``generateSweepPoints_kernel`` (``preprocess_kernel.cu:58-90``): the affine is cast to float32 and
    each output coordinate is ``((m[r,0]*x + m[r,1]*y) + m[r,2]*z) + m[r,3]`` in float32, left to right (no FMA).
    ``xyz``: (N, >=3); returns (N, 3) float32."""
    p = np.asarray(xyz, dtype=np.float32)
    m = np.asarray(T, dtype=np.float32)
    x, y, z = p[:, 0], p[:, 1], p[:, 2]
    out = np.empty((len(p), 3), np.float32)
    for r in range(3):
        out[:, r] = ((m[r, 0] * x + m[r, 1] * y) + m[r, 2] * z) + m[r, 3]
    return out


# ----------------------------------------------------------------------------------------------- rotations

def rotation_from_rpy(roll: float, pitch: float, yaw: float) -> np.ndarray:
    """tf2 ``setRPY``: ``R = Rz(yaw) @ Ry(pitch) @ Rx(roll)`` (float64)."""
    cr, sr, cp, sp, cy, sy = (math.cos(roll), math.sin(roll), math.cos(pitch), math.sin(pitch), math.cos(yaw),
                              math.sin(yaw))
    return np.array([[cy * cp, cy * sp * sr - sy * cr, cy * sp * cr + sy * sr],
                     [sy * cp, sy * sp * sr + cy * cr, sy * sp * cr - cy * sr],
                     [-sp, cp * sr, cp * cr]], dtype=np.float64)


def transform_from_xyz_rpy(x: float, y: float, z: float, roll: float = 0.0, pitch: float = 0.0,
                           yaw: float = 0.0) -> np.ndarray:
    """A (4, 4) float64 transform from a translation and tf2 RPY angles (Autoware's calibration yaml spelling)."""
    out = np.eye(4)
    out[:3, :3] = rotation_from_rpy(roll, pitch, yaw)
    out[:3, 3] = (x, y, z)
    return out


def quaternion_wxyz_from_yaw(yaw: float) -> np.ndarray:
    """``autoware_utils::create_quaternion_from_yaw``: rotation about +z, as (w, x, y, z) float64."""
    return np.array([math.cos(yaw / 2.0), 0.0, 0.0, math.sin(yaw / 2.0)], dtype=np.float64)


def yaw_from_quaternion_wxyz(q: Sequence[float]) -> float:
    """``tf2::getYaw`` of a (w, x, y, z) quaternion."""
    w, x, y, z = (float(v) for v in q)
    return math.atan2(2.0 * (w * z + x * y), 1.0 - 2.0 * (y * y + z * z))


def yaw_from_rotation(r: np.ndarray) -> float:
    """Heading of a rotation matrix (rotation of +x about +z, ``atan2(R[1,0], R[0,0])``)."""
    r = np.asarray(r, dtype=np.float64)
    return math.atan2(r[1, 0], r[0, 0])


def wrap_angle(a: Any) -> Any:
    """Wrap angles to [-pi, pi)."""
    return (np.asarray(a, dtype=np.float64) + math.pi) % (2.0 * math.pi) - math.pi


def mmdet_yaw_to_ros(yaw_net: Any) -> np.ndarray:
    """``ros_utils.cpp:55``: ``const float yaw = -box3d.yaw - pi / 2`` (float input, double arithmetic, float
    result). Returns float32 radians, not wrapped (as Autoware)."""
    y = np.asarray(yaw_net, dtype=np.float32).astype(np.float64)
    return (-y - math.pi / 2.0).astype(np.float32)


def ros_yaw_to_mmdet(yaw_ros: Any) -> np.ndarray:
    """Inverse of :func:`mmdet_yaw_to_ros` (float64): ``yaw_net = -yaw_ros - pi/2``."""
    return -np.asarray(yaw_ros, dtype=np.float64) - math.pi / 2.0