"""GNHPR dual-antenna baseline conventions and SO(3) interpolation."""

from __future__ import annotations

from dataclasses import dataclass

import numpy as np
from scipy.spatial.transform import Rotation, Slerp

@dataclass(frozen=True)
class GnhprConvention:
    """Interpretation of the GNHPR ANT1-to-ANT2 baseline.

    Heading is clockwise from north. In this vehicle the RTK ``+X`` axis is
    the baseline from the main/left antenna (ANT1) to the secondary/right
    antenna (ANT2), so it points to vehicle right rather than vehicle forward.
    GNHPR pitch is the elevation of that baseline. A two-antenna receiver
    cannot observe rotation about the baseline; the reported roll is therefore
    not treated as a third independent attitude measurement.
    """

    name: str
    heading_sign: float = -1.0
    pitch_sign: float = -1.0
    roll_sign: float = 1.0


EXPECTED_GNHPR = GnhprConvention("ant1_to_ant2__north_cw__elevation_up")
GNHPR_CANDIDATES = (
    EXPECTED_GNHPR,
    GnhprConvention("north_cw__pitch_opposite", -1.0, 1.0, 1.0),
    GnhprConvention("heading_opposite__pitch_nose_up", 1.0, -1.0, 1.0),
    GnhprConvention("heading_opposite__pitch_opposite", 1.0, 1.0, 1.0),
)


def gnhpr_to_rotation_enu_rtk(
    heading_deg: np.ndarray,
    pitch_deg: np.ndarray,
    roll_deg: np.ndarray,
    convention: GnhprConvention = EXPECTED_GNHPR,
) -> np.ndarray:
    """Build a zero-roll mathematical completion of the baseline frame.

    Only the first column (the ANT1-to-ANT2 unit vector) is physically observed
    by GNHPR. The remaining columns are a convenient gauge completion and must
    not be used as a measured full vehicle attitude.
    """

    heading = np.asarray(heading_deg, dtype=float).reshape(-1)
    pitch = np.asarray(pitch_deg, dtype=float).reshape(-1)
    roll = np.asarray(roll_deg, dtype=float).reshape(-1)
    if not (heading.size == pitch.size == roll.size):
        raise ValueError("heading, pitch and roll must have equal length")
    yaw_rad = np.deg2rad(90.0 + convention.heading_sign * heading)
    pitch_rad = np.deg2rad(convention.pitch_sign * pitch)
    # A dual-antenna baseline has no independent roll observation. Keep the
    # argument for wire-format compatibility, but never inject it into SO(3).
    roll_rad = np.zeros_like(roll)
    angles = np.column_stack([yaw_rad, pitch_rad, roll_rad])
    return Rotation.from_euler('ZYX', angles).as_matrix()


def gnhpr_to_baseline_enu(
    heading_deg: np.ndarray,
    pitch_deg: np.ndarray,
    convention: GnhprConvention = EXPECTED_GNHPR,
) -> np.ndarray:
    """Return ANT1-to-ANT2 unit vectors expressed in ENU.

    The result is independent of the unobservable GNHPR roll field.
    """

    heading = np.asarray(heading_deg, dtype=float).reshape(-1)
    pitch = np.asarray(pitch_deg, dtype=float).reshape(-1)
    if heading.size != pitch.size:
        raise ValueError("heading and pitch must have equal length")
    azimuth = np.deg2rad(heading)
    elevation = np.deg2rad(-convention.pitch_sign * pitch)
    horizontal = np.cos(elevation)
    east = np.sin(azimuth) * horizontal
    north = np.cos(azimuth) * horizontal
    up = np.sin(elevation)
    return np.column_stack([east, north, up])


def interpolate_rotations(
    source_t_s: np.ndarray,
    rotations: np.ndarray,
    query_t_s: np.ndarray,
) -> np.ndarray:
    """Slerp a monotonic SO(3) series without extrapolation."""

    source_t = np.asarray(source_t_s, dtype=float).reshape(-1)
    query_t = np.asarray(query_t_s, dtype=float).reshape(-1)
    matrices = np.asarray(rotations, dtype=float).reshape(-1, 3, 3)
    if source_t.size < 2 or matrices.shape[0] != source_t.size:
        raise ValueError("need at least two timestamped rotations")
    if np.any(np.diff(source_t) <= 0):
        unique_t, unique_indices = np.unique(source_t, return_index=True)
        source_t = unique_t
        matrices = matrices[unique_indices]
    if np.any(query_t < source_t[0]) or np.any(query_t > source_t[-1]):
        raise ValueError("rotation interpolation does not extrapolate")
    return Slerp(source_t, Rotation.from_matrix(matrices))(query_t).as_matrix()