"""Earth geometry and gravity models."""
import numpy as np
from . import dcm
from . util import mv_prod


#: Rotation rate of Earth in rad/s.
RATE = 7.2921157e-5
#: Schuller frequency.
SF = 1.2383e-3
#: Approximate value of gravity.
G0 = 9.8
#: Semi major axis of Earth ellipsoid (or radius of Earth approximately).
R0 = 6378137.0
#: Squared eccentricity of Earth ellipsoid
E2 = 6.6943799901413e-3
#: Gravity at the equator
GE = 9.7803253359
#: Gravity at the pole.
GP = 9.8321849378

F = R0 * (1 - E2) ** 0.5 * GP / (GE * R0) - 1
#: Standard gravitational parameter for Earth
MU = 3.986004418e14


def principal_radii(lat):
    """Compute the principal radii of curvature of Earth ellipsoid.

    See [1]_ (p. 41) for the definition and formulas.

    Parameters
    ----------
    lat : array_like
        Latitude.

    Returns
    -------
    re, rn : ndarray
        Radii of curvature in East and North directions respectively.

    References
    ----------
    .. [1] P. D. Groves, "Principles of GNSS, Inertial, and Multisensor
           Integrated Navigation Systems".
    """
    sin_lat = np.sin(np.deg2rad(lat))

    x = 1 - E2 * sin_lat ** 2
    re = R0 / np.sqrt(x)
    rn = re * (1 - E2) / x

    return re, rn


def gravity(lat, alt=0):
    """Compute gravity according to a theoretical model.

    See `set_model` for the explanation of the formula used.

    Parameters
    ----------
    lat : array_like
        Latitude.
    alt : array_like, optional
        Altitude. Default is 0.

    Returns
    -------
    g : ndarray
        Magnitude of the gravity.
    """
    sin_lat = np.sin(np.deg2rad(lat))
    alt = np.asarray(alt)

    return GE * (1 + F * sin_lat**2) / (1 - E2 * sin_lat**2)**0.5 * (1 - 2 * alt / R0)


def gravitation_ecef(lla):
    """Compute a vector of the gravitational force in ECEF frame.

    It accounts only for Earth mass attraction. It is computed from
    `gravity` model by eliminating the centrifugal force.

    Parameters
    ----------
    lla : array_like, shape (3,) or (n, 3)
        Latitude, longitude and altitude.

    Returns
    -------
    g0_e: ndarray, shape (3,) or (n, 3)
        Vectors of the gravitational force expressed in ECEF frame.
    """
    lat, lon, alt = np.asarray(lla).T

    sin_lat = np.sin(np.deg2rad(lat))
    cos_lat = np.cos(np.deg2rad(lat))

    re, _ = principal_radii(lat)
    rp = (re + alt) * cos_lat

    g0_g = np.zeros((3,) + sin_lat.shape)
    g0_g[1] = RATE**2 * rp * sin_lat
    g0_g[2] = -gravity(lat, alt) - RATE ** 2 * rp * cos_lat
    g0_g = g0_g.T

    Ceg = dcm.from_ll(lat, lon)

    return mv_prod(Ceg, g0_g)

@@ -2616,8 +2616,909 @@
 , g0_g)%0A
+%0A%0Adef curvature_matrix(lat, alt):%0A    %22%22%22Compute Earth curvature matrix.%0A%0A    Curvature matrix %60F%60 links linear displacement and angular rotation of%0A    ENU frame as %60rotation_n = F @ translation_n%60, where %60translation_n%60 is%0A    linear translation in ENU and %60rotation_n%60 is the corresponding small%0A    rotation vector of ENU.%0A%0A    For example %60transport_rate_n = F @ velocity_n%60.%0A%0A    Parameters%0A    ----------%0A    lat : array_like, shape (n,)%0A        Latitude.%0A    alt : array_like, shape (n,)%0A        Altitude.%0A%0A    Returns%0A    -------%0A    F: ndarray, shape (n, 3, 3)%0A        Curvature matrix.%0A    %22%22%22%0A    re, rn = principal_radii(lat)%0A    n = 1 if re.ndim == 0 else len(re)%0A%0A    result = np.zeros((n, 3, 3))%0A    result%5B:, 0, 1%5D = -1 / (rn + alt)%0A    result%5B:, 1, 0%5D = 1 / (re + alt)%0A    result%5B:, 2, 0%5D = result%5B:, 1, 0%5D * np.tan(np.deg2rad(lat))%0A%0A    return result%5B0%5D if re.ndim == 0 else result%0A

Implement earth.curvature_matrix