#!/usr/bin/env python3
"""RTK dual-antenna attitude helpers shared by prepare and SLAM delivery."""

from __future__ import annotations

import math
import re

import numpy as np

# GNHPR / UNIHEADINGA pitch is baseline elevation (far antenna higher ⇒ +pitch).
# Build the baseline-frame attitude first, then apply the fixed body yaw offset:
#   R_W_body = Rz(yaw_raw) Ry(-pitch) Rx(roll) Rz(-heading_offset)
# so pitch/roll stay about the physical baseline, even when delivering vehicle-forward.


def heading_to_enu_yaw(raw_heading_deg: float, heading_offset_deg: float = 0.0) -> tuple[float, float]:
    """Convert clockwise-from-north heading to mathematical ENU yaw (rad)."""
    corrected_heading = (raw_heading_deg + heading_offset_deg) % 360.0
    return corrected_heading, math.radians(90.0 - corrected_heading)


def _rz(yaw_rad: float) -> np.ndarray:
    c, s = math.cos(yaw_rad), math.sin(yaw_rad)
    return np.array([[c, -s, 0.0], [s, c, 0.0], [0.0, 0.0, 1.0]], dtype=float)


def _ry(pitch_rad: float) -> np.ndarray:
    c, s = math.cos(pitch_rad), math.sin(pitch_rad)
    return np.array([[c, 0.0, s], [0.0, 1.0, 0.0], [-s, 0.0, c]], dtype=float)


def _rx(roll_rad: float) -> np.ndarray:
    c, s = math.cos(roll_rad), math.sin(roll_rad)
    return np.array([[1.0, 0.0, 0.0], [0.0, c, -s], [0.0, s, c]], dtype=float)


def attitude_rotation(
    yaw_rad: float,
    pitch_deg: float = 0.0,
    roll_deg: float = 0.0,
) -> np.ndarray:
    """ENU←baseline rotation: Rz(yaw) Ry(-pitch) Rx(roll).

    Positive ``pitch_deg`` elevates baseline X (slave higher than master).
    """
    return _rz(float(yaw_rad)) @ _ry(-math.radians(float(pitch_deg))) @ _rx(math.radians(float(roll_deg)))


def rtk_body_rotation(
    raw_heading_deg: float,
    heading_offset_deg: float = 0.0,
    pitch_deg: float = 0.0,
    roll_deg: float = 0.0,
) -> np.ndarray:
    """ENU←delivered RTK body frame.

    Pitch/roll are applied in the raw baseline frame; ``heading_offset_deg`` then
    rotates that frame into the delivered body (0 = baseline X, -90 = vehicle
    forward when baseline points vehicle-right on this vehicle).
    """
    _, yaw_baseline = heading_to_enu_yaw(raw_heading_deg, 0.0)
    return attitude_rotation(yaw_baseline, pitch_deg, roll_deg) @ _rz(-math.radians(float(heading_offset_deg)))


def rotation_to_quat_xyzw(rotation: np.ndarray) -> np.ndarray:
    r = np.asarray(rotation, dtype=float)
    tr = float(np.trace(r))
    if tr > 0.0:
        s = math.sqrt(tr + 1.0) * 2.0
        q = np.array(
            [(r[2, 1] - r[1, 2]) / s, (r[0, 2] - r[2, 0]) / s, (r[1, 0] - r[0, 1]) / s, 0.25 * s],
            dtype=float,
        )
    else:
        i = int(np.argmax(np.diag(r)))
        if i == 0:
            s = math.sqrt(1.0 + r[0, 0] - r[1, 1] - r[2, 2]) * 2.0
            q = np.array(
                [0.25 * s, (r[0, 1] + r[1, 0]) / s, (r[0, 2] + r[2, 0]) / s, (r[2, 1] - r[1, 2]) / s],
                dtype=float,
            )
        elif i == 1:
            s = math.sqrt(1.0 + r[1, 1] - r[0, 0] - r[2, 2]) * 2.0
            q = np.array(
                [(r[0, 1] + r[1, 0]) / s, 0.25 * s, (r[1, 2] + r[2, 1]) / s, (r[0, 2] - r[2, 0]) / s],
                dtype=float,
            )
        else:
            s = math.sqrt(1.0 + r[2, 2] - r[0, 0] - r[1, 1]) * 2.0
            q = np.array(
                [(r[0, 2] + r[2, 0]) / s, (r[1, 2] + r[2, 1]) / s, 0.25 * s, (r[1, 0] - r[0, 1]) / s],
                dtype=float,
            )
    if q[3] < 0.0:
        q = -q
    return q / np.linalg.norm(q)


def parse_pitch_roll_from_heading_raw(raw_utf8: bytes | str | None) -> tuple[float | None, float | None]:
    """Best-effort pitch/roll from a stored GNHPR/UNIHEADINGA raw line."""
    if raw_utf8 is None:
        return None, None
    text = raw_utf8.decode("ascii", "ignore") if isinstance(raw_utf8, (bytes, bytearray)) else str(raw_utf8)
    text = text.strip()
    if "GNHPR" in text:
        parts = text.split(",")
        if len(parts) >= 5:
            try:
                return float(parts[3]), float(parts[4])
            except ValueError:
                return None, None
    if "UNIHEADINGA" in text.upper() or "HEADINGA" in text.upper():
        payload = text.split(";", 1)[-1]
        fields = payload.split(",")
        if len(fields) >= 5:
            try:
                return float(fields[4]), 0.0
            except ValueError:
                return None, None
    match = re.search(r",(-?\d+(?:\.\d+)?),(-?\d+(?:\.\d+)?),\d,", text)
    if match:
        try:
            return float(match.group(1)), float(match.group(2))
        except ValueError:
            return None, None
    return None, None