"""Small WGS84 geodesy helpers used by the RTK--IMU calibration path."""

from __future__ import annotations

import numpy as np


WGS84_A_M = 6378137.0
WGS84_F = 1.0 / 298.257223563
WGS84_E2 = WGS84_F * (2.0 - WGS84_F)


def geodetic_to_ecef(
    latitude_deg: np.ndarray,
    longitude_deg: np.ndarray,
    altitude_m: np.ndarray,
) -> np.ndarray:
    """Convert WGS84 latitude/longitude/ellipsoidal height to ECEF metres."""

    latitude = np.deg2rad(np.asarray(latitude_deg, dtype=float))
    longitude = np.deg2rad(np.asarray(longitude_deg, dtype=float))
    altitude = np.asarray(altitude_m, dtype=float)
    latitude, longitude, altitude = np.broadcast_arrays(latitude, longitude, altitude)
    sin_lat = np.sin(latitude)
    cos_lat = np.cos(latitude)
    radius = WGS84_A_M / np.sqrt(1.0 - WGS84_E2 * sin_lat**2)
    x = (radius + altitude) * cos_lat * np.cos(longitude)
    y = (radius + altitude) * cos_lat * np.sin(longitude)
    z = (radius * (1.0 - WGS84_E2) + altitude) * sin_lat
    return np.stack([x, y, z], axis=-1)


def geodetic_to_enu(
    latitude_deg: np.ndarray,
    longitude_deg: np.ndarray,
    altitude_m: np.ndarray,
    *,
    origin_latitude_deg: float | None = None,
    origin_longitude_deg: float | None = None,
    origin_altitude_m: float | None = None,
) -> tuple[np.ndarray, tuple[float, float, float]]:
    """Convert WGS84 samples to a local east/north/up frame.

    When no origin is supplied, the first finite sample is used. The returned
    origin tuple is ``(latitude_deg, longitude_deg, altitude_m)``.
    """

    lat = np.asarray(latitude_deg, dtype=float).reshape(-1)
    lon = np.asarray(longitude_deg, dtype=float).reshape(-1)
    alt = np.asarray(altitude_m, dtype=float).reshape(-1)
    if not (lat.size == lon.size == alt.size):
        raise ValueError("latitude, longitude and altitude must have equal length")
    finite = np.isfinite(lat) & np.isfinite(lon) & np.isfinite(alt)
    if not np.any(finite):
        raise ValueError("no finite geodetic sample")
    first = int(np.flatnonzero(finite)[0])
    lat0 = float(lat[first] if origin_latitude_deg is None else origin_latitude_deg)
    lon0 = float(lon[first] if origin_longitude_deg is None else origin_longitude_deg)
    alt0 = float(alt[first] if origin_altitude_m is None else origin_altitude_m)

    ecef = geodetic_to_ecef(lat, lon, alt)
    ecef0 = geodetic_to_ecef(np.array(lat0), np.array(lon0), np.array(alt0)).reshape(3)
    delta = ecef - ecef0
    phi = np.deg2rad(lat0)
    lam = np.deg2rad(lon0)
    rotation = np.array(
        [
            [-np.sin(lam), np.cos(lam), 0.0],
            [-np.sin(phi) * np.cos(lam), -np.sin(phi) * np.sin(lam), np.cos(phi)],
            [np.cos(phi) * np.cos(lam), np.cos(phi) * np.sin(lam), np.sin(phi)],
        ],
        dtype=float,
    )
    return delta @ rotation.T, (lat0, lon0, alt0)