"""Transform XYZ (ECEF) locations <=> LLA (Lat, Lon, Alt) locations.

These functions translate between XYZ in meters and LLA in degrees / meters. They
take and return numpy arrays.

# Get the XYZ in meters from LLA in degrees and meters
lla_to_xyz(lla_ddm)

# Get the LLA in degrees and meters from XYZ in meters
xyz_to_lla(xyz_m)
"""

from __future__ import annotations

from typing import TYPE_CHECKING

import numpy as np

from gri_utils import constants

if TYPE_CHECKING:
    from collections.abc import Sequence

    from numpy.typing import NDArray


def lla_to_xyz(lla_ddm: NDArray[np.floating] | Sequence) -> NDArray[np.floating]:
    """Convert Lat, Lon, Alt in degrees and meters to XYZ ECEF in meters.

    Args:
        lla_ddm(np.ndarray|Sequence): Lat, Lon, Alt in degrees and meters. Shape (3,)
            for single point or (N, 3) for N points.

    Returns:
        XYZ ECEF in meters as a numpy array. Shape (3,) or (N, 3) matching input.
    """
    lla_ddm = np.asarray(lla_ddm)
    lat_rad = np.radians(lla_ddm[..., 0])
    lon_rad = np.radians(lla_ddm[..., 1])
    alt_m = lla_ddm[..., 2]

    cos_lon = np.cos(lon_rad)
    sin_lon = np.sin(lon_rad)
    cos_lat = np.cos(lat_rad)
    sin_lat = np.sin(lat_rad)

    nlat = constants.ER_M / (1 - constants.ECCENTRICITY_SQUARED * sin_lat**2) ** 0.5
    x = (nlat + alt_m) * cos_lat * cos_lon
    y = (nlat + alt_m) * cos_lat * sin_lon
    z = ((1 - constants.ECCENTRICITY_SQUARED) * nlat + alt_m) * sin_lat
    return np.stack((x, y, z), axis=-1)


def xyz_to_lla(xyz_m: NDArray[np.floating] | Sequence) -> NDArray[np.floating]:
    """Convert XYZ in meters to Lat, Lon, Alt in degrees and meters.

    Args:
        xyz_m(np.ndarray|Sequence): XYZ ECEF coordinate in meters. Shape (3,) for single
            point or (N, 3) for N points.

    Returns:
        Lat, Lon, Alt in degrees and meters as a numpy array. Shape (3,) or (N, 3)
            matching input.
    """
    xyz_m = np.asarray(xyz_m)
    p = (xyz_m[..., 0] ** 2 + xyz_m[..., 1] ** 2) ** 0.5
    theta = np.arctan2(constants.ER_M * xyz_m[..., 2], constants.PR_M * p)
    stheta = np.sin(theta)
    ctheta = np.cos(theta)
    lon = np.arctan2(xyz_m[..., 1], xyz_m[..., 0])
    lat = np.arctan2(
        xyz_m[..., 2]
        + (constants.ER_M**2 / constants.PR_M - constants.PR_M) * stheta**3,
        p - (constants.ER_M - constants.PR_M**2 / constants.ER_M) * ctheta**3,
    )
    n = constants.ER_M / (1 - constants.ECCENTRICITY_SQUARED * np.sin(lat) ** 2) ** 0.5

    # Standard altitude formula has a singularity at the poles where cos(lat) -> 0.
    # Near poles (p < 1m), use simplified polar geometry: alt = |z| - polar_radius.
    # The 1m threshold ensures sub-millimeter accuracy at the transition.
    near_pole = p < 1.0
    alt_standard = p / np.cos(lat) - n
    alt_polar = np.abs(xyz_m[..., 2]) - constants.PR_M
    alt = np.where(near_pole, alt_polar, alt_standard)

    return np.stack(
        (
            np.degrees(lat),
            np.degrees(lon),
            alt,
        ),
        axis=-1,
    )
