"""
aruco_math.py — Pure-math helpers for ArUco detection (Issue #627).

No ROS2 dependencies — importable in unit tests without a ROS2 install.

Provides:
  rot_mat_to_quat(R)         → (qx, qy, qz, qw)
  rvec_to_quat(rvec)         → (qx, qy, qz, qw)   [via cv2.Rodrigues]
  tvec_distance(tvec)        → float   (Euclidean norm)
  tvec_yaw_rad(tvec)         → float   (atan2(tx, tz) — lateral bearing error)
  MarkerPose dataclass       — per-marker detection result
"""

from __future__ import annotations

import math
from dataclasses import dataclass, field
from typing import List, Optional, Tuple

import numpy as np


# ── Rotation helpers ──────────────────────────────────────────────────────────

def rot_mat_to_quat(R: np.ndarray) -> Tuple[float, float, float, float]:
    """
    Convert a 3×3 rotation matrix to quaternion (qx, qy, qz, qw).

    Uses Shepperd's method for numerical stability.
    R must be a valid rotation matrix (‖R‖₂ = 1, det R = +1).
    """
    R = np.asarray(R, dtype=np.float64)
    trace = R[0, 0] + R[1, 1] + R[2, 2]

    if trace > 0.0:
        s = 0.5 / math.sqrt(trace + 1.0)
        w = 0.25 / s
        x = (R[2, 1] - R[1, 2]) * s
        y = (R[0, 2] - R[2, 0]) * s
        z = (R[1, 0] - R[0, 1]) * s
    elif R[0, 0] > R[1, 1] and R[0, 0] > R[2, 2]:
        s = 2.0 * math.sqrt(1.0 + R[0, 0] - R[1, 1] - R[2, 2])
        w = (R[2, 1] - R[1, 2]) / s
        x = 0.25 * s
        y = (R[0, 1] + R[1, 0]) / s
        z = (R[0, 2] + R[2, 0]) / s
    elif R[1, 1] > R[2, 2]:
        s = 2.0 * math.sqrt(1.0 + R[1, 1] - R[0, 0] - R[2, 2])
        w = (R[0, 2] - R[2, 0]) / s
        x = (R[0, 1] + R[1, 0]) / s
        y = 0.25 * s
        z = (R[1, 2] + R[2, 1]) / s
    else:
        s = 2.0 * math.sqrt(1.0 + R[2, 2] - R[0, 0] - R[1, 1])
        w = (R[1, 0] - R[0, 1]) / s
        x = (R[0, 2] + R[2, 0]) / s
        y = (R[1, 2] + R[2, 1]) / s
        z = 0.25 * s

    # Normalise
    norm = math.sqrt(x * x + y * y + z * z + w * w)
    if norm < 1e-10:
        return 0.0, 0.0, 0.0, 1.0
    inv = 1.0 / norm
    return x * inv, y * inv, z * inv, w * inv


def rvec_to_quat(rvec: np.ndarray) -> Tuple[float, float, float, float]:
    """
    Convert a Rodrigues rotation vector (3,) or (1,3) to quaternion.

    Requires cv2 for cv2.Rodrigues(); falls back to manual exponential map
    if cv2 is unavailable.
    """
    vec = np.asarray(rvec, dtype=np.float64).ravel()
    try:
        import cv2
        R, _ = cv2.Rodrigues(vec)
    except ImportError:
        # Manual Rodrigues exponential map
        angle = float(np.linalg.norm(vec))
        if angle < 1e-12:
            return 0.0, 0.0, 0.0, 1.0
        axis = vec / angle
        K = np.array([
            [     0.0, -axis[2],  axis[1]],
            [ axis[2],      0.0, -axis[0]],
            [-axis[1],  axis[0],      0.0],
        ])
        R = np.eye(3) + math.sin(angle) * K + (1.0 - math.cos(angle)) * K @ K

    return rot_mat_to_quat(R)


# ── Metric helpers ────────────────────────────────────────────────────────────

def tvec_distance(tvec: np.ndarray) -> float:
    """Euclidean distance from camera to marker (m)."""
    t = np.asarray(tvec, dtype=np.float64).ravel()
    return float(math.sqrt(t[0] ** 2 + t[1] ** 2 + t[2] ** 2))


def tvec_yaw_rad(tvec: np.ndarray) -> float:
    """
    Horizontal bearing error to marker (rad).

    atan2(tx, tz): +ve when marker is to the right of the camera optical axis.
    Zero when the marker is directly in front.
    """
    t = np.asarray(tvec, dtype=np.float64).ravel()
    return float(math.atan2(t[0], t[2]))


# ── Per-marker result ─────────────────────────────────────────────────────────

@dataclass
class MarkerPose:
    """Pose result for a single detected ArUco marker."""
    marker_id:  int
    tvec:       np.ndarray    # (3,) translation in camera optical frame (m)
    rvec:       np.ndarray    # (3,) Rodrigues rotation vector
    corners:    np.ndarray    # (4, 2) image-plane corner coordinates (px)

    # Derived (computed lazily)
    _distance_m: Optional[float] = field(default=None, repr=False)
    _yaw_rad:    Optional[float] = field(default=None, repr=False)
    _quat:       Optional[Tuple[float, float, float, float]] = field(default=None, repr=False)

    @property
    def distance_m(self) -> float:
        if self._distance_m is None:
            self._distance_m = tvec_distance(self.tvec)
        return self._distance_m

    @property
    def yaw_rad(self) -> float:
        if self._yaw_rad is None:
            self._yaw_rad = tvec_yaw_rad(self.tvec)
        return self._yaw_rad

    @property
    def lateral_m(self) -> float:
        """Lateral offset (m); +ve = marker is to the right."""
        return float(self.tvec[0])

    @property
    def forward_m(self) -> float:
        """Forward distance (Z component in camera frame, m)."""
        return float(self.tvec[2])

    @property
    def quat(self) -> Tuple[float, float, float, float]:
        """Quaternion (qx, qy, qz, qw) from rvec."""
        if self._quat is None:
            self._quat = rvec_to_quat(self.rvec)
        return self._quat
