"""SE(3) helpers backed by GTSAM `Pose3` (E-CC-HELPERS / AZ-264 / AZ-277).

Implements the `se3_utils` contract v1.0.0 at
`_docs/02_document/contracts/shared_helpers/se3_utils.md`. Stateless,
pure functions; strict caller-orthogonalisation invariant.
"""

from __future__ import annotations

from typing import Final

import gtsam
import numpy as np

__all__ = [
    "SE3",
    "Se3InvalidMatrixError",
    "adjoint",
    "exp_map",
    "is_valid_rotation",
    "log_map",
    "matrix_to_se3",
    "se3_to_matrix",
]


SE3 = gtsam.Pose3

# Tolerance for the orthogonality / determinant / bottom-row checks. Tight
# enough that an ill-conditioned rotation from a real consumer is caught;
# loose enough to not reject within-noise GTSAM output. Documented at the
# contract level for symmetry with consumer tests.
_DEFAULT_ROT_ATOL: Final[float] = 1e-6
_DEFAULT_BOTTOM_ROW: Final[np.ndarray] = np.array([0.0, 0.0, 0.0, 1.0], dtype=np.float64)

# Small-angle Taylor cutoff for `exp_map` stability (AC-3). Below this twist
# norm we delegate to GTSAM's first-order Taylor fallback rather than risk
# the `sin(theta)/theta` numerator under-flowing to zero.
_SMALL_ANGLE_THRESHOLD: Final[float] = 1e-10


class Se3InvalidMatrixError(ValueError):
    """Raised when an input matrix violates the SE(3) shape / dtype / orthogonality contract."""


def _require_float64(array: np.ndarray, *, name: str) -> None:
    if array.dtype != np.float64:
        raise Se3InvalidMatrixError(
            f"{name}: helpers operate strictly on dtype=float64; got {array.dtype}"
        )


def _require_shape(array: np.ndarray, expected: tuple[int, ...], *, name: str) -> None:
    if array.shape != expected:
        raise Se3InvalidMatrixError(f"{name}: expected shape {expected}; got {array.shape}")


def is_valid_rotation(R_3x3: np.ndarray, *, atol: float = _DEFAULT_ROT_ATOL) -> bool:
    """Return True iff `R_3x3` is an orthogonal rotation with positive determinant."""
    if not isinstance(R_3x3, np.ndarray):
        return False
    if R_3x3.shape != (3, 3) or R_3x3.dtype != np.float64:
        return False
    drift = R_3x3.T @ R_3x3 - np.eye(3, dtype=np.float64)
    if np.linalg.norm(drift, ord="fro") > atol:
        return False
    if np.linalg.det(R_3x3) < 0:
        return False
    return True


def matrix_to_se3(T_4x4: np.ndarray, *, atol: float = _DEFAULT_ROT_ATOL) -> SE3:
    """Convert a 4x4 homogeneous-transform matrix into a GTSAM `Pose3`.

    Strict orthogonality contract: callers MUST pre-orthogonalise their
    rotation matrices. Non-orthogonal inputs, negative-determinant
    rotations, malformed bottom rows, and `float32` inputs all raise
    `Se3InvalidMatrixError` — the helper never silently re-orthogonalises.
    """
    if not isinstance(T_4x4, np.ndarray):
        raise Se3InvalidMatrixError(
            f"matrix_to_se3: expected np.ndarray; got {type(T_4x4).__name__}"
        )
    _require_shape(T_4x4, (4, 4), name="matrix_to_se3")
    _require_float64(T_4x4, name="matrix_to_se3")

    bottom_row = T_4x4[3]
    if not np.array_equal(bottom_row, _DEFAULT_BOTTOM_ROW):
        raise Se3InvalidMatrixError(
            f"matrix_to_se3: bottom row must be [0, 0, 0, 1]; got {bottom_row.tolist()}"
        )

    R = T_4x4[:3, :3]
    drift = R.T @ R - np.eye(3, dtype=np.float64)
    drift_norm = float(np.linalg.norm(drift, ord="fro"))
    if drift_norm > atol:
        raise Se3InvalidMatrixError(
            f"matrix_to_se3: rotation is not orthogonal "
            f"(||R^T R - I||_F = {drift_norm:.3e} > atol={atol:.1e}); "
            f"caller must orthogonalise before invoking the helper"
        )

    det = float(np.linalg.det(R))
    if det < 0:
        raise Se3InvalidMatrixError(
            f"matrix_to_se3: rotation has negative determinant ({det:.3e}); "
            f"mirror rotations are not valid SE(3) members"
        )

    return SE3(T_4x4)


def se3_to_matrix(pose: SE3) -> np.ndarray:
    """Return the 4x4 homogeneous matrix for `pose` as `float64`."""
    return np.ascontiguousarray(pose.matrix(), dtype=np.float64)


def exp_map(xi: np.ndarray) -> SE3:
    """Exponential map: se(3) twist (6,) -> SE(3) pose.

    Near-identity inputs (twist norm below the small-angle threshold)
    fall back to the identity pose rather than relying on the
    full-precision `sin(theta)/theta` expansion.
    """
    if not isinstance(xi, np.ndarray):
        raise Se3InvalidMatrixError(f"exp_map: expected np.ndarray; got {type(xi).__name__}")
    _require_shape(xi, (6,), name="exp_map")
    _require_float64(xi, name="exp_map")

    if float(np.linalg.norm(xi)) < _SMALL_ANGLE_THRESHOLD:
        return SE3()
    return SE3.Expmap(xi)


def log_map(pose: SE3) -> np.ndarray:
    """Logarithm map: SE(3) pose -> se(3) twist (6,)."""
    return np.ascontiguousarray(SE3.Logmap(pose), dtype=np.float64)


def adjoint(pose: SE3) -> np.ndarray:
    """Adjoint matrix (6x6) of `pose` for body-frame -> world-frame twist transport."""
    return np.ascontiguousarray(pose.AdjointMap(), dtype=np.float64)